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__`.

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,16 @@
+{# ----------------------------------------------------------------------------
+ # 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" %}
+from ._init import *
+{% endif %}

symforce/codegen/cam_package_codegen.py

@@ -293,7 +293,8 @@ def generate(config: CodegenConfig, output_dir: str = None) -> str:
                 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: str = None) -> str:
         )
 
         # 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

@@ -143,8 +143,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",
@@ -163,7 +165,7 @@ def test_codegen_python(self) -> None:
             actual_dir=output_dir, expected_dir=os.path.join(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", os.path.join(output_dir, "sym", "__init__.py")
@@ -200,9 +202,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,
@@ -218,6 +222,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
@@ -243,10 +269,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)
 
     def test_matrix_indexing_python(self) -> None:
         """
@@ -270,9 +298,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,
@@ -502,7 +532,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]]:
@@ -514,11 +543,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())
@@ -557,14 +588,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++
@@ -1072,7 +1103,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
         )
 
@@ -1083,7 +1114,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:
@@ -1140,11 +1171,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

@@ -70,7 +70,7 @@ def gen_code(self, output_dir: str) -> 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,
         )
@@ -81,7 +81,7 @@ def gen_code(self, output_dir: str) -> 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

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