memoize: bugfix for functions and handling for classes

core/ipyflow/data_model/symbol.py

@@ -1246,7 +1246,7 @@ def make_memoize_comparable_for_obj(
                 if size > cls._MAX_MEMOIZE_COMPARABLE_SIZE:
                     return cls.NULL, None, -1
                 comparable.append(inner_comp)
-            return type(obj)(comparable), cls._equal, size
+            return comparable, cls._equal, size
         else:
             # hacks to check if they are arrays or dataframes without explicitly importing these
             module = getattr(type(obj), "__module__", "")
@@ -1262,8 +1262,8 @@ def make_memoize_comparable_for_obj(
     def make_memoize_comparable(
         self,
     ) -> Tuple[Any, Optional[Callable[[Any, Any], bool]]]:
-        if isinstance(self.stmt_node, ast.FunctionDef):
-            return astunparse.unparse(self.stmt_node)
+        if isinstance(self.stmt_node, (ast.ClassDef, ast.FunctionDef)):
+            return astunparse.unparse(self.stmt_node), self._equal
         obj, eq, size = self.make_memoize_comparable_for_obj(self.obj)
         if size > self._MAX_MEMOIZE_COMPARABLE_SIZE:
             return self.NULL, None

core/test/test_memoization.py

@@ -98,3 +98,30 @@ def test_sets():
     assert flow().global_scope["y"].obj == {0, 1}
     assert second.is_memoized
     assert second.skipped_due_to_memoization_ctr > 0
+
+
+def test_functions():
+    first = cells(run_cell("def foo(): return 42"))
+    second = cells(run_cell("%%memoize\ny = foo() + 1", cell_id="second"))
+    assert shell().user_ns["y"] == 43
+    assert flow().global_scope["y"].obj == 43
+    assert second.is_memoized
+    assert second.skipped_due_to_memoization_ctr == -1
+    run_cell("def foo(): return 42", cell_id=first.id)
+    second = cells(run_cell("%%memoize\ny = foo() + 1", cell_id="second"))
+    assert shell().user_ns["y"] == 43
+    assert flow().global_scope["y"].obj == 43
+    assert second.is_memoized
+    assert second.skipped_due_to_memoization_ctr > 0
+    run_cell("def foo(): return 44", cell_id=first.id)
+    second = cells(run_cell("%%memoize\ny = foo() + 1", cell_id="second"))
+    assert shell().user_ns["y"] == 45
+    assert flow().global_scope["y"].obj == 45
+    assert second.is_memoized
+    assert second.skipped_due_to_memoization_ctr == -1
+    run_cell("def foo(): return 42", cell_id=first.id)
+    second = cells(run_cell("%%memoize\ny = foo() + 1", cell_id="second"))
+    assert shell().user_ns["y"] == 43
+    assert flow().global_scope["y"].obj == 43
+    assert second.is_memoized
+    assert second.skipped_due_to_memoization_ctr > 0