add test

marimo/_runtime/editor/_editor.py

@@ -1,6 +1,7 @@
+# Copyright 2024 Marimo. All rights reserved.
 from marimo._ast.visitor import VariableData
-from marimo._data.get_datasets import get_table_manager_or_none
 from marimo._output.rich_help import mddoc
+from marimo._plugins.ui._impl.tables.utils import get_table_manager_or_none
 from marimo._runtime.context.types import (
     ContextNotInitializedError,
     get_context,
@@ -9,9 +10,15 @@
 
 @mddoc
 def register_datasource(obj: object, name: str) -> None:
-    """Register a datasource with the current context.
+    """Register a datasource.
 
-    This datasource will be available to other cells in the same context.
+    This registered object will be available in the global scope of the
+    notebook, including as a variable in the graph.
+
+    WARNING: This function may cause unintended bugs in reactivity, since
+    defined variables cannot be statically analyzed. Also, this can be
+    confusing for users if used inappropriately to flood the global scope.
+    Please be mindful of this function.
 
     **Args:**
 
@@ -23,6 +30,9 @@ def register_datasource(obj: object, name: str) -> None:
     except ContextNotInitializedError:
         return
 
+    if ctx.execution_context is None:
+        return
+
     if get_table_manager_or_none(obj) is None:
         raise ValueError(f"Failed to get table data for variable {name}")
 
@@ -33,6 +43,6 @@ def register_datasource(obj: object, name: str) -> None:
     cell.defs.add(name)
     cell.variable_data[name] = [VariableData("variable")]
     if name in ctx.graph.definitions:
-        ctx.graph.definitions[name].append(cell_id)
+        ctx.graph.definitions[name].add(cell_id)
     else:
-        ctx.graph.definitions.update({name: [cell_id]})
+        ctx.graph.definitions.update({name: {cell_id}})

tests/_runtime/editor/test_editor.py

@@ -0,0 +1,58 @@
+import sys
+import types
+
+import pytest
+
+from marimo._dependencies.dependencies import DependencyManager
+from marimo._runtime.runtime import Kernel
+from tests.conftest import ExecReqProvider
+
+registering_fn_module = types.ModuleType("registering_fn_module")
+exec(
+    """
+import marimo as mo
+import pandas as pd
+
+def registering_fn() -> None:
+    print('hi')
+    df = pd.DataFrame({'a': [1], 'b': [2]})
+    name_fn = lambda x: x
+    mo.editor.register_datasource(df, name_fn('test_var_name'))
+""",
+    registering_fn_module.__dict__,
+)
+
+# Add this module to `sys.modules`
+sys.modules["registering_fn_module"] = registering_fn_module
+
+
+class TestCellRun:
+    @staticmethod
+    @pytest.mark.skipif(
+        condition=not DependencyManager.pandas.has(),
+        reason="requires matplotlib",
+    )
+    async def test_register_datasource(
+        execution_kernel: Kernel, exec_req: ExecReqProvider
+    ) -> None:
+        k = execution_kernel
+        await k.run(
+            [
+                exec_req.get(
+                    """
+                    import registering_fn_module
+                    registering_fn_module.registering_fn()
+                    """
+                )
+            ]
+        )
+
+        assert k.globals["registering_fn_module"]
+        assert (
+            "test_var_name" in k.globals
+        ), "test_var_name not found in globals."
+        import pandas as pd
+
+        assert isinstance(
+            k.globals["test_var_name"], pd.DataFrame
+        ), "test_var_name is not a pandas DataFrame."