Make sym and other genned pkgs namespace packages

notebooks/tutorials/codegen_tutorial.ipynb

@@ -390,16 +390,16 @@
     "params.L = [0.5, 0.3]\n",
     "params.m = [0.3, 0.2]\n",
     "\n",
-    "gen_module = codegen_util.load_generated_package(\n",
-    "    namespace, double_pendulum_python_data.function_dir\n",
+    "gen_double_pendulum = codegen_util.load_generated_function(\n",
+    "    \"double_pendulum\", double_pendulum_python_data.function_dir\n",
     ")\n",
-    "gen_module.double_pendulum(ang, dang, consts, params)"
+    "gen_double_pendulum(ang, dang, consts, params)"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -413,7 +413,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.9"
+   "version": "3.8.14"
   }
  },
  "nbformat": 4,

symforce/codegen/backends/python/python_config.py

@@ -33,13 +33,16 @@ class PythonConfig(CodegenConfig):
                    times.
         reshape_vectors: Allow rank 1 ndarrays to be passed in for row and column vectors by
                          automatically reshaping the input.
+        namespace_package: Generate the package as a namespace package, meaning it can be split
+                           across multiple directories.
     """
 
     doc_comment_line_prefix: str = ""
     line_length: int = 100
     use_eigen_types: bool = True
     use_numba: bool = False
     reshape_vectors: bool = True
+    namespace_package: bool = True
 
     @classmethod
     def backend_name(cls) -> str:
@@ -50,10 +53,11 @@ def template_dir(cls) -> Path:
         return CURRENT_DIR / "templates"
 
     def templates_to_render(self, generated_file_name: str) -> T.List[T.Tuple[str, str]]:
-        return [
-            ("function/FUNCTION.py.jinja", f"{generated_file_name}.py"),
-            ("function/__init__.py.jinja", "__init__.py"),
-        ]
+        templates = [("function/FUNCTION.py.jinja", f"{generated_file_name}.py")]
+        if self.namespace_package:
+            return templates + [("function/namespace_init.py.jinja", "__init__.py")]
+        else:
+            return templates + [("function/__init__.py.jinja", "__init__.py")]
 
     def printer(self) -> CodePrinter:
         return python_code_printer.PythonCodePrinter()

symforce/codegen/backends/python/templates/function/__init__.py.jinja

@@ -2,4 +2,3 @@
  # SymForce - Copyright 2022, Skydio, Inc.
  # This source code is under the Apache 2.0 license found in the LICENSE file.
  # ---------------------------------------------------------------------------- #}
-from .{{ spec.name }} import {{ spec.name }}

symforce/codegen/backends/python/templates/function/namespace_init.py.jinja

@@ -0,0 +1,20 @@
+{# ----------------------------------------------------------------------------
+ # SymForce - Copyright 2022, Skydio, Inc.
+ # This source code is under the Apache 2.0 license found in the LICENSE file.
+ # ---------------------------------------------------------------------------- #}
+{% if pkg_namespace == "sym" %}
+"""
+Python runtime geometry package.
+"""
+
+{% endif %}
+# Make package a namespace package by adding other portions to the __path__
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore[has-type]
+# https://github.com/python/mypy/issues/1422
+{% if pkg_namespace == "sym" %}
+# Import sym types into the top level of the package. Needs to be done in all portions of
+# the package since namespace packages do not evaluate every copy of `__init__.py`, just
+# the first discovered. Import must occur after __path__ has been modified to ensure the
+# sym._init module can be located.
+from ._init import *
+{% endif %}

symforce/codegen/cam_package_codegen.py

@@ -293,7 +293,8 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
                 all_types=list(geo_package_codegen.DEFAULT_GEO_TYPES) + list(DEFAULT_CAM_TYPES),
                 numeric_epsilon=sf.numeric_epsilon,
             ),
-            output_path=cam_package_dir / "__init__.py",
+            output_path=cam_package_dir
+            / ("_init.py" if config.namespace_package else "__init__.py"),
         )
 
         for name in ("cam_package_python_test.py",):

symforce/codegen/codegen.py

@@ -484,7 +484,7 @@ def generate_function(
         # Namespace of this function + generated types
         self.namespace = namespace
 
-        template_data = dict(self.common_data(), spec=self)
+        template_data = dict(self.common_data(), spec=self, pkg_namespace=namespace)
         template_dir = self.config.template_dir()
 
         backend_name = self.config.backend_name()

symforce/codegen/geo_package_codegen.py

@@ -217,14 +217,20 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
         )
 
         # Package init
+        if config.namespace_package:
+            templates.add(
+                template_path=Path("function", "namespace_init.py.jinja"),
+                data=dict(pkg_namespace="sym"),
+                output_path=package_dir / "__init__.py",
+            )
         templates.add(
             template_path=Path("geo_package", "__init__.py.jinja"),
             data=dict(
                 Codegen.common_data(),
                 all_types=DEFAULT_GEO_TYPES,
                 numeric_epsilon=sf.numeric_epsilon,
             ),
-            output_path=package_dir / "__init__.py",
+            output_path=package_dir / ("_init.py" if config.namespace_package else "__init__.py"),
         )
 
         # Test example

symforce/opt/numeric_factor.py

@@ -75,10 +75,7 @@ def from_file_python(
         """
         assert all(opt_key in keys for opt_key in optimized_keys)
         function_dir = Path(output_dir) / "python" / "symforce" / namespace
-        linearization_function = getattr(
-            codegen_util.load_generated_package(f"{namespace}.{name}", function_dir),
-            name,
-        )
+        linearization_function = codegen_util.load_generated_function(name, function_dir)
         return cls(
             keys=keys, optimized_keys=optimized_keys, linearization_function=linearization_function
         )

test/symforce_codegen_test.py

@@ -144,8 +144,10 @@ def test_codegen_python(self) -> None:
         """
         inputs, outputs = self.build_values()
 
+        config = codegen.PythonConfig(namespace_package=False)
+
         python_func = codegen.Codegen(
-            inputs=inputs, outputs=outputs, config=codegen.PythonConfig(), name="python_function"
+            inputs=inputs, outputs=outputs, config=config, name="python_function"
         )
         shared_types = {
             "values_vec": "values_vec_t",
@@ -164,7 +166,7 @@ def test_codegen_python(self) -> None:
             actual_dir=output_dir, expected_dir=TEST_DATA_DIR / (namespace + "_data")
         )
 
-        geo_package_codegen.generate(config=codegen.PythonConfig(), output_dir=output_dir)
+        geo_package_codegen.generate(config=config, output_dir=output_dir)
 
         geo_pkg = codegen_util.load_generated_package("sym", output_dir / "sym" / "__init__.py")
         values_vec_t = codegen_util.load_generated_lcmtype(
@@ -199,9 +201,11 @@ def test_codegen_python(self) -> None:
 
         big_matrix = np.zeros((5, 5))
 
-        gen_module = codegen_util.load_generated_package(namespace, codegen_data.function_dir)
+        python_function = codegen_util.load_generated_function(
+            "python_function", codegen_data.function_dir
+        )
         # TODO(nathan): Split this test into several different functions
-        (foo, bar, scalar_vec_out, values_vec_out, values_vec_2D_out) = gen_module.python_function(
+        (foo, bar, scalar_vec_out, values_vec_out, values_vec_2D_out) = python_function(
             x,
             y,
             rot,
@@ -217,6 +221,28 @@ def test_codegen_python(self) -> None:
         self.assertStorageNear(foo, x ** 2 + rot.data[3])
         self.assertStorageNear(bar, constants.epsilon + sf.sin(y) + x ** 2)
 
+    def test_return_geo_type_from_generated_python_function(self) -> None:
+        """
+        Tests that the function (returning a Rot3) generated by codegen.Codegen.generate_function()
+        with the default PythonConfig can be called.
+        When test was created, if you tried to do this, the error:
+            AttributeError: module 'sym' has no attribute 'Rot3'
+        would be raised.
+        """
+
+        def identity() -> sf.Rot3:
+            return sf.Rot3.identity()
+
+        output_dir = self.make_output_dir("sf_test_return_geo_type_from_generated_python_function")
+
+        codegen_data = codegen.Codegen.function(
+            func=identity, config=codegen.PythonConfig()
+        ).generate_function(output_dir=output_dir)
+
+        gen_identity = codegen_util.load_generated_function("identity", codegen_data.function_dir)
+
+        gen_identity()
+
     def test_matrix_order_python(self) -> None:
         """
         Tests that codegen.Codegen.generate_function() renders matrices correctly
@@ -242,10 +268,12 @@ def matrix_order() -> sf.M23:
             func=matrix_order, config=codegen.PythonConfig()
         ).generate_function(namespace=namespace, output_dir=output_dir)
 
-        pkg = codegen_util.load_generated_package(namespace, codegen_data.function_dir)
+        gen_matrix_order = codegen_util.load_generated_function(
+            "matrix_order", codegen_data.function_dir
+        )
 
-        self.assertEqual(pkg.matrix_order().shape, m23.SHAPE)
-        self.assertStorageNear(pkg.matrix_order(), m23)
+        self.assertEqual(gen_matrix_order().shape, m23.SHAPE)
+        self.assertStorageNear(gen_matrix_order(), m23)
 
     @unittest.skipIf(importlib.util.find_spec("numba") is None, "Requires numba")
     def test_matrix_indexing_python(self) -> None:
@@ -271,9 +299,11 @@ def gen_pass_matrices(use_numba: bool, reshape_vectors: bool) -> T.Any:
                 output_names=["row_out", "col_out", "mat_out"],
             ).generate_function(namespace=namespace, output_dir=output_dir)
 
-            pkg = codegen_util.load_generated_package(namespace, generated_files.function_dir)
+            genned_func = codegen_util.load_generated_function(
+                "pass_matrices", generated_files.function_dir
+            )
 
-            return pkg.pass_matrices
+            return genned_func
 
         def assert_config_works(
             use_numba: bool,
@@ -503,7 +533,6 @@ def test_sparse_output_python(self) -> None:
         argument of codegen.Codegen.__init__ is set appropriately.
         """
         output_dir = self.make_output_dir("sf_test_sparse_output_python")
-        namespace = "sparse_output_python"
         x, y, z = sf.symbols("x y z")
 
         def matrix_output(x: sf.Scalar, y: sf.Scalar, z: sf.Scalar) -> T.List[T.List[sf.Scalar]]:
@@ -515,11 +544,13 @@ def matrix_output(x: sf.Scalar, y: sf.Scalar, z: sf.Scalar) -> T.List[T.List[sf.
             name="sparse_output_func",
             config=codegen.PythonConfig(),
             sparse_matrices=["out"],
-        ).generate_function(namespace=namespace, output_dir=output_dir)
+        ).generate_function(namespace="sparse_output_python", output_dir=output_dir)
 
-        pkg = codegen_util.load_generated_package(namespace, codegen_data.function_dir)
+        sparse_output_func = codegen_util.load_generated_function(
+            "sparse_output_func", codegen_data.function_dir
+        )
 
-        output = pkg.sparse_output_func(1, 2, 3)
+        output = sparse_output_func(1, 2, 3)
 
         self.assertIsInstance(output, sparse.csc_matrix)
         self.assertTrue((output.todense() == matrix_output(1, 2, 3)).all())
@@ -558,14 +589,14 @@ def numba_test_func(x: sf.V3) -> sf.V2:
         output_function = numba_test_func_codegen_data.function_dir / "numba_test_func.py"
         self.compare_or_update_file(expected_code_file, output_function)
 
-        gen_module = codegen_util.load_generated_package(
-            "sym", numba_test_func_codegen_data.function_dir
+        gen_func = codegen_util.load_generated_function(
+            "numba_test_func", numba_test_func_codegen_data.function_dir
         )
 
         x = np.array([1, 2, 3])
-        y = gen_module.numba_test_func(x)
+        y = gen_func(x)
         self.assertTrue((y == np.array([[1, 2]]).T).all())
-        self.assertTrue(hasattr(gen_module.numba_test_func, "__numba__"))
+        self.assertTrue(hasattr(gen_func, "__numba__"))
 
     # -------------------------------------------------------------------------
     # C++
@@ -1071,7 +1102,7 @@ def test_function_dataclass(dataclass: TestDataclass1, x: sf.Scalar) -> sf.V3:
             func=test_function_dataclass, config=codegen.PythonConfig()
         )
         dataclass_codegen_data = dataclass_codegen.generate_function()
-        gen_module = codegen_util.load_generated_package(
+        gen_func = codegen_util.load_generated_function(
             "test_function_dataclass", dataclass_codegen_data.function_dir
         )
 
@@ -1082,7 +1113,7 @@ def test_function_dataclass(dataclass: TestDataclass1, x: sf.Scalar) -> sf.V3:
         dataclass_t.v2.v0 = 1
 
         # make sure it runs
-        gen_module.test_function_dataclass(dataclass_t, 1)
+        gen_func(dataclass_t, 1)
 
     @slow_on_sympy
     def test_function_explicit_template_instantiation(self) -> None:
@@ -1139,11 +1170,11 @@ class MyDataclass:
         )
 
         # Make sure it runs
-        gen_module = codegen_util.load_generated_package(namespace, codegen_data.function_dir)
+        gen_func = codegen_util.load_generated_function(name, codegen_data.function_dir)
         my_dataclass_t = codegen_util.load_generated_lcmtype(
             namespace, "my_dataclass_t", codegen_data.python_types_dir
         )()
-        return_rot = gen_module.codegen_dataclass_in_values_test(my_dataclass_t)
+        return_rot = gen_func(my_dataclass_t)
         self.assertEqual(return_rot.data, my_dataclass_t.rot.data)
 
 

test/symforce_databuffer_codegen_test.py

@@ -69,7 +69,7 @@ def gen_code(self, output_dir: Path) -> None:
         )
 
         # Also test that the generated python code runs
-        gen_module = codegen_util.load_generated_package(
+        buffer_func = codegen_util.load_generated_function(
             "buffer_func",
             py_codegen_data.function_dir,
         )
@@ -80,7 +80,7 @@ def gen_code(self, output_dir: Path) -> None:
         # 2 * buffer[b^2 - a^2] + (a+b)
         # 2 * buffer[3] + 3
         expected = 9
-        result_numeric = gen_module.buffer_func(buffer_numeric, a_numeric, b_numeric)
+        result_numeric = buffer_func(buffer_numeric, a_numeric, b_numeric)
 
         self.assertStorageNear(expected, result_numeric)
 

test/symforce_function_codegen_test_data/symengine/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py

@@ -1,7 +1,9 @@
 # -----------------------------------------------------------------------------
 # This file was autogenerated by symforce from template:
-#     function/__init__.py.jinja
+#     function/namespace_init.py.jinja
 # Do NOT modify by hand.
 # -----------------------------------------------------------------------------
 
-from .codegen_dataclass_in_values_test import codegen_dataclass_in_values_test
+# Make package a namespace package by adding other portions to the __path__
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore[has-type]
+# https://github.com/python/mypy/issues/1422