Merge pull request #3

Enable module imports and better environment variable handling

dbt_py/main.py

@@ -15,9 +15,7 @@
 
 PROJECT_ROOT = pathlib.Path(__file__).parent.parent
 # Python-style ref, e.g. `package.module.submodule`
-PACKAGE_ROOT: str = os.environ.get("DBT_PY_PACKAGE_ROOT", "custom")
-# The name to associate with the package
-PACKAGE_NAME: str = os.environ.get("DBT_PY_PACKAGE_NAME", PACKAGE_ROOT)
+PACKAGE_ROOT: str = os.environ.get("DBT_PY_PACKAGE_ROOT")
 
 
 def import_submodules(
@@ -30,6 +28,10 @@ def import_submodules(
     - https://stackoverflow.com/a/25562415/10730311
     """
     package = importlib.import_module(package_name)
+    if not hasattr(package, "__path__"):
+        # `package` is a module, don't recurse any further
+        return {}
+
     results = {}
     for loader, name, is_pkg in pkgutil.walk_packages(package.__path__):
         full_name = f"{package.__name__}.{name}"
@@ -44,9 +46,12 @@ def new_get_context_modules() -> dict[str, dict[str, Any]]:
     """
     Append the custom modules into the whitelisted dbt modules.
     """
-    import_submodules(PACKAGE_ROOT)
+    package_root: str = PACKAGE_ROOT or "custom"
+    package_name: str = os.environ.get("DBT_PY_PACKAGE_NAME") or package_root
+
+    import_submodules(package_root)
     modules = _get_context_modules()
-    modules[PACKAGE_NAME] = importlib.import_module(PACKAGE_ROOT)  # type: ignore
+    modules[package_name] = importlib.import_module(package_root)  # type: ignore
 
     return modules
 

tests/integration/jaffle-shop/dbt-commands.sh

@@ -1,6 +1,6 @@
 # Just for reference
 # $env:DBT_PY_PACKAGE_ROOT = "tests.integration.jaffle-shop.dbt_py_test"
-# $env:DBT_PY_PACKAGE_NAME = "dbt_py_test"
+# $env:DBT_PY_PACKAGE_NAME = "custom_py"
 dbt    clean   --project-dir tests/integration/jaffle-shop --profiles-dir tests/integration/jaffle-shop --no-clean-project-files-only
 dbt-py compile --project-dir tests/integration/jaffle-shop --profiles-dir tests/integration/jaffle-shop
 rm dbt_packages

tests/integration/jaffle-shop/models/example.sql

@@ -1,5 +1,5 @@
-{{ modules.dbt_py_test.some_module.select_final() }}
-{{ modules.dbt_py_test.some_module.salutation("World") }}
+{{ modules.custom_py.some_module.select_final() }}
+{{ modules.custom_py.some_module.salutation("World") }}
 
-{{ modules.dbt_py_test.select_final() }}
-{{ modules.dbt_py_test.salutation("World") }}
+{{ modules.custom_py.select_final() }}
+{{ modules.custom_py.salutation("World") }}

tests/integration/test__integration.py

@@ -37,7 +37,7 @@ def mock_env(monkeypatch) -> None:
     monkeypatch.setenv(
         "DBT_PY_PACKAGE_ROOT", "tests.integration.jaffle-shop.dbt_py_test"
     )
-    monkeypatch.setenv("DBT_PY_PACKAGE_NAME", "dbt_py_test")
+    monkeypatch.setenv("DBT_PY_PACKAGE_NAME", "custom_py")
 
 
 def test__integration(mock_env) -> None: