From 9e6f3a6ed22183d4c693fc73fddd57b624a70d37 Mon Sep 17 00:00:00 2001
From: "Jason R. Coombs" <jaraco@jaraco.com>
Date: Tue, 11 Jun 2024 16:41:25 -0400
Subject: [PATCH] Add support for redirecting for a namespace package.

---
 src/editables/redirector.py | 41 ++++++++++++++++++++++++++++++++-----
 tests/test_redirects.py     | 20 ++++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/src/editables/redirector.py b/src/editables/redirector.py
index 7bdef59..70ada27 100644
--- a/src/editables/redirector.py
+++ b/src/editables/redirector.py
@@ -3,26 +3,57 @@
 import importlib.util
 import sys
 from types import ModuleType
-from typing import Dict, Optional, Sequence, Union
+from typing import Dict, Optional, Sequence, Set, Union
 
 ModulePath = Optional[Sequence[Union[bytes, str]]]
 
 
 class RedirectingFinder(importlib.abc.MetaPathFinder):
     _redirections: Dict[str, str] = {}
+    _parents: Set[str] = set()
 
     @classmethod
     def map_module(cls, name: str, path: str) -> None:
         cls._redirections[name] = path
+        cls._parents.update(cls.parents(name))
+
+    @classmethod
+    def parents(cls, name):
+        """
+        Given a full name, generate all parents.
+
+        >>> list(RedirectingFinder.parents('a.b.c.d'))
+        ['a.b.c', 'a.b', 'a']
+        """
+        base, sep, name = name.rpartition('.')
+        if base:
+            yield base
+            yield from cls.parents(base)
 
     @classmethod
     def find_spec(
         cls, fullname: str, path: ModulePath = None, target: Optional[ModuleType] = None
     ) -> Optional[importlib.machinery.ModuleSpec]:
-        if "." in fullname:
-            return None
-        if path is not None:
-            return None
+        return cls.spec_from_parent(fullname) or cls.spec_from_redirect(fullname)
+
+    @classmethod
+    def spec_from_parent(
+        cls, fullname: str
+    ) -> Optional[importlib.machinery.ModuleSpec]:
+        if fullname in cls._parents:
+            return importlib.util.spec_from_loader(
+                fullname,
+                importlib.machinery.NamespaceLoader(
+                    fullname,
+                    path=[],
+                    path_finder=cls.find_spec,
+                ),
+            )
+
+    @classmethod
+    def spec_from_redirect(
+        cls, fullname: str
+    ) -> Optional[importlib.machinery.ModuleSpec]:
         try:
             redir = cls._redirections[fullname]
         except KeyError:
diff --git a/tests/test_redirects.py b/tests/test_redirects.py
index 65a2059..5e37440 100644
--- a/tests/test_redirects.py
+++ b/tests/test_redirects.py
@@ -84,6 +84,26 @@ def test_redirects(tmp_path):
         assert pkg.sub.val == 42
 
 
+def test_namespace_redirects(tmp_path):
+    project = tmp_path / "project"
+    project_files = {
+        "ns.pkg": {
+            "__init__.py": "val = 42",
+            "sub.py": "val = 42",
+        }
+    }
+    build(project, project_files)
+
+    with save_import_state():
+        F.install()
+        F.map_module("ns.pkg", project / "ns.pkg" / "__init__.py")
+
+        import ns.pkg.sub
+
+        assert ns.pkg.val == 42
+        assert ns.pkg.sub.val == 42
+
+
 def test_cache_invalidation():
     F.install()
     # assert that the finder matches importlib's expectations