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

Author: bradley-solliday-skydio

gen/python/sym/__init__.py

@@ -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,

gen/python/sym/_init.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()

notebooks/tutorials/codegen_tutorial.ipynb

@@ -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/python_config.py

@@ -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/backends/python/templates/function/__init__.py.jinja

@@ -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/backends/python/templates/function/namespace_init.py.jinja

@@ -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/cam_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/codegen/codegen.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
         )