wip state - type checking and naive packing start - still broken mess

Author: RonnyPfannschmidt

src/execnet/_dictimport.py

@@ -5,6 +5,7 @@
 import base64
 import json
 import os
+import pkgutil
 import sys
 import types
 import zlib
@@ -13,13 +14,16 @@
 from importlib.metadata import Distribution
 from importlib.metadata import DistributionFinder
 from typing import IO
-from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
+from typing import NamedTuple
+from typing import NewType
+from typing import Protocol
 from typing import Sequence
+from typing import cast
+from typing import runtime_checkable
 
-if TYPE_CHECKING:
-    pass
+ModuleName = NewType("ModuleName", str)
 
 
 class DictDistribution(Distribution):
@@ -35,56 +39,58 @@ def locate_file(self, path: str | os.PathLike[str]) -> os.PathLike[str]:
         raise FileNotFoundError(path)
 
 
+class ModuleInfo(NamedTuple):
+    name: ModuleName
+    is_pkg: bool
+    source: str
+
+
 class DictImporter(DistributionFinder, Loader):
     """a limited loader/importer for distributins send via json-lines"""
 
-    def __init__(self, sources: dict[str, str], distribution: DictDistribution):
+    def __init__(
+        self, sources: dict[ModuleName, ModuleInfo], distribution: DictDistribution
+    ):
         self.sources = sources
         self.distribution = distribution
 
     def find_distributions(
-        self, context: DistributionFinder.Context = DistributionFinder.Context()
+        self, context: DistributionFinder.Context | None = None
     ) -> Iterable[Distribution]:
+        # TODO: filter
         return [self.distribution]
 
     def find_module(
         self, fullname: str, path: Sequence[str | bytes] | None = None
     ) -> Loader | None:
-        if fullname in self.sources:
-            return self
-        if fullname + ".__init__" in self.sources:
+        if ModuleName(fullname) in self.sources:
             return self
         return None
 
-    def load_module(self, fullname):
+    def load_module(self, fullname: str) -> types.ModuleType:
         # print "load_module:",  fullname
-        from types import ModuleType
 
-        try:
-            s = self.sources[fullname]
-            is_pkg = False
-        except KeyError:
-            s = self.sources[fullname + ".__init__"]
-            is_pkg = True
+        info = self.sources[ModuleName(fullname)]
 
-        co = compile(s, fullname, "exec")
-        module = sys.modules.setdefault(fullname, ModuleType(fullname))
+        co = compile(info.source, fullname, "exec")
+        module = sys.modules.setdefault(fullname, types.ModuleType(fullname))
         module.__loader__ = self
-        if is_pkg:
+        if info.is_pkg:
             module.__path__ = [fullname]
 
         exec(co, module.__dict__)
         return sys.modules[fullname]
 
     def get_source(self, name: str) -> str | None:
-        res = self.sources.get(name)
+        res = self.sources.get(ModuleName(name))
         if res is None:
-            res = self.sources.get(name + ".__init__")
-        return res
+            return None
+        else:
+            return res.source
 
 
 def bootstrap(
-    modules: dict[str, str],
+    modules: dict[ModuleName, ModuleInfo],
     distribution: dict[str, str],
     entry: str,
     args: dict[str, Any],
@@ -100,7 +106,7 @@ def bootstrap(
     entry_func(**args)
 
 
-def bootstrap_stdin(stream: IO) -> None:
+def bootstrap_stdin(stream: IO[bytes] | IO[str]) -> None:
     bootstrap_args = decode_b85_zip_json(stream.readline())
     bootstrap(**bootstrap_args)
 
@@ -111,9 +117,26 @@ def decode_b85_zip_json(encoded: bytes | str):
     return json.loads(unpacked)
 
 
-def naive_pack_module(module: types.ModuleType, dist: Distribution):
+@runtime_checkable
+class SourceProvidingLoader(Protocol):
+    def get_source(self, name: str) -> str:
+        ...
+
+
+def naive_pack_module(module: types.ModuleType, dist: Distribution) -> object:
     assert module.__file__ is not None
     assert module.__path__
+    data: dict[ModuleName, ModuleInfo] = {}
+    for info in pkgutil.walk_packages(module.__path__, f"{module.__name__}."):
+        spec = info.module_finder.find_spec(info.name, None)
+        assert spec is not None
+        loader = cast(SourceProvidingLoader, spec.loader)
+
+        source = loader.get_source(info.name)
+        data[ModuleName(info.name)] = ModuleInfo(
+            name=ModuleName(info.name), is_pkg=info.ispkg, source=source
+        )
+    return data
 
 
 if __name__ == "__main__":