From 5e6ce26628319042ae18bad8add4ffc8d776ad6f Mon Sep 17 00:00:00 2001
From: bradley-solliday-skydio <bradley.solliday@skydio.com>
Date: Mon, 31 Oct 2022 18:57:06 -0700
Subject: [PATCH] Make sym and other genned pkgs namespace packages

Namespace packages are packages whose `__path__` (which tells the python
import system to look for sub-packages and sub-modules) has multiple
directories, meaning it can have portions spread out across multiple
directories.

Previously, `sym` and the other generated packages were not namespace
packages. This caused issues when generated python packages in the `sym`
namespace attempted to access, for example, `sym.Rot3` (which they might
do if the generated function returns a `sym.Rot3`). Since the generated
package itself was named `sym`, but `Rot3` was not defined locally, an
`AttributeError` would be raised.

While an alternative would have been to instead not use the name `sym`
for generated packages (using, say, `sym_gen` instead), for reasons I
didn't fully look into, we still want to generate our code within the
`sym` namespace (generating into `sym.gen` was considered as an option
but to do so would require the changes in this commit regardless).

While currently we haven't needed to generate any functions returning a
`sym` class in a package named `sym`, we intend to soon add a `sym.util`
package which will be used in a lot of places. That can't be done until
this namespace conflict is resolved.

Note, while a python2 compatible namespace package has multiple
`__init__.py` files for the top-level package spread across different
locations, only one of them will be executed. This makes it difficult to
re-export the contents of sub-modules into the top-level namespace.

The normal way to re-export a name is to write
``` python3
from .sub_module import name
```
However, since the sub-modules can be created dynamically, it is
impossible to re-export all names in this manner, as the first
`__init__.py` that is created has no way of knowing what names one might
want to re-export from subsequent modules.

It is possible to put all names one wishes to export in a standard file,
say `_init.py`, then dynamically search for such files and execute their
contents, but we considered the additional complexity to be too large of
a burden (as users would have a harder time understand their generated
code, and this would give future maintainers a hard time).

And so, we decided to simply stop re-exporting any names in the
`__init__.py`'s of generated code (kind of in the style of pep 420
python3 packages).

This makes loading a generated function more difficult if one uses
`codegen_util.load_generated_package`, as now simply importing a
generated package won't give you access to any of the package's
contents. However, this is what `codegen_util.load_generated_function`
is for, so hopefully the user experience shouldn't be too negatively
impacted.

The one exception to the general ban of re-exporting names is the `sym`
package, as we still wish to be able to do
``` python3
import sym

sym.Rot3.identity()
```
However, because all sub-modules we wish to export from the `sym`
package are known at code-gen time, allowing this is not difficult.
This only applies to names in the core `sym` package, and any additional
user generated code in the `sym` package will not be re-rexported in the
top-level namespace.

A user can prevent their package from being generated as a namespace
package by setting the `namespace_package` field of their `PythonConfig`
to `False`. This is useful in our testing as it is the generated code
being tested that is imported, not, for example, the checked in `sym`
package code which is being imported.

As a last note, I believe `pkgutil.extend_path` only checks for portions
of the package on the `sys.path`, and doesn't check for any portions
than can only be found by finders on the `sys.meta_path` (for example,
`symforce` itself is found by a finder on the `sys.meta_path` but not on
the `sys.path` during an editable pip install). I don't expect this
lapse to pose a problem, and addressing it immediately might just make
the `__init__.py`s look more complicated than they need to be, but if
this does become a problem, know that the situation can be partially
addressed trying to find the spec using the finders, and adding the
spec's `submodule_search_locations` if found to the `__path__`.
---
 gen/python/sym/__init__.py                    | 21 +++---
 gen/python/sym/_init.py                       | 21 ++++++
 notebooks/tutorials/codegen_tutorial.ipynb    | 10 +--
 .../codegen/backends/python/python_config.py  | 12 ++-
 .../templates/function/__init__.py.jinja      |  1 -
 .../function/namespace_init.py.jinja          | 20 +++++
 symforce/codegen/cam_package_codegen.py       |  3 +-
 symforce/codegen/codegen.py                   |  2 +-
 symforce/codegen/geo_package_codegen.py       |  8 +-
 symforce/opt/numeric_factor.py                |  5 +-
 test/symforce_codegen_test.py                 | 73 +++++++++++++------
 test/symforce_databuffer_codegen_test.py      |  4 +-
 .../python/symforce/codegen_test/__init__.py  |  6 +-
 .../symforce/codegen_python_test/__init__.py  |  2 -
 .../python/symforce/buffer_test/__init__.py   |  6 +-
 .../python/symforce/codegen_test/__init__.py  |  6 +-
 .../symforce/codegen_python_test/__init__.py  |  2 -
 .../python/symforce/buffer_test/__init__.py   |  6 +-
 18 files changed, 144 insertions(+), 64 deletions(-)
 create mode 100644 gen/python/sym/_init.py
 create mode 100644 symforce/codegen/backends/python/templates/function/namespace_init.py.jinja

diff --git a/gen/python/sym/__init__.py b/gen/python/sym/__init__.py
index de8c975ef..0f403da1a 100644
--- a/gen/python/sym/__init__.py
+++ b/gen/python/sym/__init__.py
@@ -1,21 +1,18 @@
 # -----------------------------------------------------------------------------
 # This file was autogenerated by symforce from template:
-#     geo_package/__init__.py.jinja
+#     function/namespace_init.py.jinja
 # Do NOT modify by hand.
 # -----------------------------------------------------------------------------
 
 """
 Python runtime geometry package.
 """
-from .atan_camera_cal import ATANCameraCal
-from .double_sphere_camera_cal import DoubleSphereCameraCal
-from .equirectangular_camera_cal import EquirectangularCameraCal
-from .linear_camera_cal import LinearCameraCal
-from .polynomial_camera_cal import PolynomialCameraCal
-from .pose2 import Pose2
-from .pose3 import Pose3
-from .rot2 import Rot2
-from .rot3 import Rot3
-from .spherical_camera_cal import SphericalCameraCal
 
-epsilon = 2.220446049250313e-15
+# 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
+# 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 *
diff --git a/gen/python/sym/_init.py b/gen/python/sym/_init.py
new file mode 100644
index 000000000..de8c975ef
--- /dev/null
+++ b/gen/python/sym/_init.py
@@ -0,0 +1,21 @@
+# -----------------------------------------------------------------------------
+# This file was autogenerated by symforce from template:
+#     geo_package/__init__.py.jinja
+# Do NOT modify by hand.
+# -----------------------------------------------------------------------------
+
+"""
+Python runtime geometry package.
+"""
+from .atan_camera_cal import ATANCameraCal
+from .double_sphere_camera_cal import DoubleSphereCameraCal
+from .equirectangular_camera_cal import EquirectangularCameraCal
+from .linear_camera_cal import LinearCameraCal
+from .polynomial_camera_cal import PolynomialCameraCal
+from .pose2 import Pose2
+from .pose3 import Pose3
+from .rot2 import Rot2
+from .rot3 import Rot3
+from .spherical_camera_cal import SphericalCameraCal
+
+epsilon = 2.220446049250313e-15
diff --git a/notebooks/tutorials/codegen_tutorial.ipynb b/notebooks/tutorials/codegen_tutorial.ipynb
index 1273f2915..c8d8f5561 100644
--- a/notebooks/tutorials/codegen_tutorial.ipynb
+++ b/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,
diff --git a/symforce/codegen/backends/python/python_config.py b/symforce/codegen/backends/python/python_config.py
index 28e9f8665..a60ac433d 100644
--- a/symforce/codegen/backends/python/python_config.py
+++ b/symforce/codegen/backends/python/python_config.py
@@ -33,6 +33,8 @@ 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 = ""
@@ -40,6 +42,7 @@ class PythonConfig(CodegenConfig):
     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()
diff --git a/symforce/codegen/backends/python/templates/function/__init__.py.jinja b/symforce/codegen/backends/python/templates/function/__init__.py.jinja
index eafe9282d..e95f1ca85 100644
--- a/symforce/codegen/backends/python/templates/function/__init__.py.jinja
+++ b/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 }}
diff --git a/symforce/codegen/backends/python/templates/function/namespace_init.py.jinja b/symforce/codegen/backends/python/templates/function/namespace_init.py.jinja
new file mode 100644
index 000000000..395009b0a
--- /dev/null
+++ b/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 %}
diff --git a/symforce/codegen/cam_package_codegen.py b/symforce/codegen/cam_package_codegen.py
index e1065acf5..f20f5d1d1 100644
--- a/symforce/codegen/cam_package_codegen.py
+++ b/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",):
diff --git a/symforce/codegen/codegen.py b/symforce/codegen/codegen.py
index 1409d440d..735d773ae 100644
--- a/symforce/codegen/codegen.py
+++ b/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()
diff --git a/symforce/codegen/geo_package_codegen.py b/symforce/codegen/geo_package_codegen.py
index 91d9c3c06..7cfec35be 100644
--- a/symforce/codegen/geo_package_codegen.py
+++ b/symforce/codegen/geo_package_codegen.py
@@ -217,6 +217,12 @@ 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(
@@ -224,7 +230,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
                 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
diff --git a/symforce/opt/numeric_factor.py b/symforce/opt/numeric_factor.py
index bb883a469..32671a7ba 100644
--- a/symforce/opt/numeric_factor.py
+++ b/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
         )
diff --git a/test/symforce_codegen_test.py b/test/symforce_codegen_test.py
index 443137da4..c177bda65 100644
--- a/test/symforce_codegen_test.py
+++ b/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)
 
 
diff --git a/test/symforce_databuffer_codegen_test.py b/test/symforce_databuffer_codegen_test.py
index d896b0efd..aba7bb39f 100644
--- a/test/symforce_databuffer_codegen_test.py
+++ b/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)
 
diff --git a/test/symforce_function_codegen_test_data/symengine/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py b/test/symforce_function_codegen_test_data/symengine/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py
index 5b40fc66c..6c7d316e1 100644
--- a/test/symforce_function_codegen_test_data/symengine/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py
+++ b/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
diff --git a/test/symforce_function_codegen_test_data/symengine/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py b/test/symforce_function_codegen_test_data/symengine/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
index a3db6f52f..3148fc07b 100644
--- a/test/symforce_function_codegen_test_data/symengine/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
+++ b/test/symforce_function_codegen_test_data/symengine/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
@@ -3,5 +3,3 @@
 #     function/__init__.py.jinja
 # Do NOT modify by hand.
 # -----------------------------------------------------------------------------
-
-from .python_function import python_function
diff --git a/test/symforce_function_codegen_test_data/symengine/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py b/test/symforce_function_codegen_test_data/symengine/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py
index b7c3d4456..6c7d316e1 100644
--- a/test/symforce_function_codegen_test_data/symengine/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py
+++ b/test/symforce_function_codegen_test_data/symengine/databuffer_codegen_test_data/python/symforce/buffer_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 .buffer_func import buffer_func
+# 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
diff --git a/test/symforce_function_codegen_test_data/sympy/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py b/test/symforce_function_codegen_test_data/sympy/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py
index 5b40fc66c..6c7d316e1 100644
--- a/test/symforce_function_codegen_test_data/sympy/codegen_dataclass_in_values_test_data/python/symforce/codegen_test/__init__.py
+++ b/test/symforce_function_codegen_test_data/sympy/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
diff --git a/test/symforce_function_codegen_test_data/sympy/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py b/test/symforce_function_codegen_test_data/sympy/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
index a3db6f52f..3148fc07b 100644
--- a/test/symforce_function_codegen_test_data/sympy/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
+++ b/test/symforce_function_codegen_test_data/sympy/codegen_python_test_data/python/symforce/codegen_python_test/__init__.py
@@ -3,5 +3,3 @@
 #     function/__init__.py.jinja
 # Do NOT modify by hand.
 # -----------------------------------------------------------------------------
-
-from .python_function import python_function
diff --git a/test/symforce_function_codegen_test_data/sympy/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py b/test/symforce_function_codegen_test_data/sympy/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py
index b7c3d4456..6c7d316e1 100644
--- a/test/symforce_function_codegen_test_data/sympy/databuffer_codegen_test_data/python/symforce/buffer_test/__init__.py
+++ b/test/symforce_function_codegen_test_data/sympy/databuffer_codegen_test_data/python/symforce/buffer_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 .buffer_func import buffer_func
+# 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