omlish 0.0.0.dev404__py3-none-any.whl → 0.0.0.dev406__py3-none-any.whl

omlish/manifests/loading.py

@@ -4,16 +4,13 @@
 Should be kept somewhat lightweight - used in cli entrypoints.
 
 TODO:
- - persisted caching support - {pkg_name: manifests}
  - real relative cls names - shouldn't need parent package names
- - *require* loaded class names - special All sentinel for explicit all
-  - ! late instantiation !
  - TypeMap style weak cache of issubclass queries
   - wait.. lazily load the class for virtual subclass queries? xor support virtual bases?
+ - weakref class dict keys?
 """
 import dataclasses as dc
 import importlib.machinery
-import importlib.resources
 import importlib.util
 import json
 import os.path
@@ -23,6 +20,9 @@ import typing as ta
 from .types import Manifest
 
 
+T = ta.TypeVar('T')
+
+
 ##
 
 
@@ -31,30 +31,42 @@ class ManifestLoader:
         def __init__(
                 self,
                 package: 'ManifestLoader.LoadedPackage',
-                manifest: Manifest,
+                resolved: 'ManifestLoader._ResolvedManifest',
         ) -> None:
             super().__init__()
 
             self._package = package
-            self._manifest = manifest
+            self._resolved = resolved
+
+        def __repr__(self) -> str:
+            return (
+                f'{self.__class__.__name__}@{id(self):x}('
+                f'package={self._package.name!r}, '
+                f'module={self._resolved.module!r}, '
+                f'class_key={self._resolved.class_key!r}'
+                f')'
+            )
 
         @property
         def package(self) -> 'ManifestLoader.LoadedPackage':
             return self._package
 
+        @property
+        def module(self) -> str:
+            return self._resolved.module
+
+        @property
+        def class_key(self) -> str:
+            return self._resolved.class_key
+
         @property
         def manifest(self) -> Manifest:
-            return self._manifest
+            return self._resolved.manifest
 
         @property
         def loader(self) -> 'ManifestLoader':
             return self._package.loader
 
-        @property
-        def class_key(self) -> str:
-            [(cls_key, value_dct)] = self._manifest.value.items()
-            return cls_key
-
         _value: ta.Any
 
         def value(self) -> ta.Any:
@@ -63,7 +75,7 @@ class ManifestLoader:
             except AttributeError:
                 pass
 
-            value = self.loader._instantiate_loaded_manifest(self)  # noqa
+            value = self.loader._instantiate_resolved_manifest(self._resolved)  # noqa
             self._value = value
             return value
 
@@ -80,6 +92,9 @@ class ManifestLoader:
 
         _manifests: ta.Sequence['ManifestLoader.LoadedManifest']
 
+        def __repr__(self) -> str:
+            return f'{self.__class__.__name__}@{id(self):x}(name={self._name!r})'
+
         @property
         def loader(self) -> 'ManifestLoader':
             return self._loader
@@ -92,56 +107,206 @@ class ManifestLoader:
         def manifests(self) -> ta.Sequence['ManifestLoader.LoadedManifest']:
             return self._manifests
 
+        _manifests_by_class_key: ta.Mapping[str, ta.Sequence['ManifestLoader.LoadedManifest']]
+
+        @property
+        def manifests_by_class_key(self) -> ta.Mapping[str, ta.Sequence['ManifestLoader.LoadedManifest']]:
+            try:
+                return self._manifests_by_class_key
+            except AttributeError:
+                pass
+
+            dct: dict = {}
+            for m in self._manifests:
+                try:
+                    lst = dct[m.class_key]
+                except KeyError:
+                    lst = dct[m.class_key] = []
+                lst.append(m)
+            self._manifests_by_class_key = dct
+            return dct
+
+    ##
+
+    @dc.dataclass(frozen=True)
+    class Config:
+        package_scan_root_dirs: ta.Optional[ta.Collection[str]] = None
+
+        discover_packages: ta.Optional[bool] = None
+        discover_packages_fallback_scan_root_dirs: ta.Optional[ta.Collection[str]] = None
+
+        module_remap: ta.Optional[ta.Mapping[str, str]] = None
+
+        value_instantiator: ta.Optional[ta.Callable[..., ta.Any]] = None
+
+        def __post_init__(self) -> None:
+            if isinstance(self.package_scan_root_dirs, str):
+                raise TypeError(self.package_scan_root_dirs)
+            if isinstance(self.discover_packages_fallback_scan_root_dirs, bool):
+                raise TypeError(self.discover_packages_fallback_scan_root_dirs)
+
+        @classmethod
+        def merge(cls, *configs: 'ManifestLoader.Config') -> 'ManifestLoader.Config':
+            kw: dict = {}
+            for c in configs:
+                for k, v in dc.asdict(c).items():
+                    if v is None:
+                        continue
+                    elif k in ('package_scan_root_dirs', 'discover_packages_fallback_scan_root_dirs'):  # noqa
+                        kw[k] = [*kw.get(k, []), *v]
+                    elif k == 'module_remap':
+                        kw[k] = {**kw.get(k, {}), **v}
+                    else:
+                        kw[k] = v
+            return cls(**kw)
+
+        def __or__(self, other: 'ManifestLoader.Config') -> 'ManifestLoader.Config':
+            return ManifestLoader.Config.merge(self, other)
+
     def __init__(
             self,
-            *,
-            module_remap: ta.Optional[ta.Mapping[str, str]] = None,
-            value_instantiator: ta.Optional[ta.Callable[..., ta.Any]] = None,
+            config: Config,
     ) -> None:
         super().__init__()
 
-        self._value_instantiator = value_instantiator
-        self._module_remap = module_remap or {}
+        self._config = config
 
         self._lock = threading.RLock()
 
-        self._module_reverse_remap = {v: k for k, v in self._module_remap.items()}
+        self._module_remap = config.module_remap or {}
+        self._module_reverse_remap = {v: k for k, v in (self._module_remap or {}).items()}
 
         self._loaded_classes: ta.Dict[str, type] = {}
         self._loaded_packages: ta.Dict[str, ta.Optional[ManifestLoader.LoadedPackage]] = {}
 
         self._scanned_package_root_dirs: ta.Dict[str, ta.Sequence[str]] = {}
 
-    #
+    @property
+    def config(self) -> Config:
+        return self._config
 
     @classmethod
-    def kwargs_from_entry_point(
+    def config_from_entry_point(
             cls,
             globals: ta.Mapping[str, ta.Any],  # noqa
-            *,
-            module_remap: ta.Optional[ta.Mapping[str, str]] = None,
-            **kwargs: ta.Any,
-    ) -> ta.Dict[str, ta.Any]:
+    ) -> Config:
         rm: ta.Dict[str, str] = {}
 
-        if module_remap:
-            rm.update(module_remap)
-
         if '__name__' in globals and '__spec__' in globals:
             name: str = globals['__name__']
             spec: importlib.machinery.ModuleSpec = globals['__spec__']
             if '__main__' not in rm and name == '__main__':
                 rm[spec.name] = '__main__'
 
-        return dict(module_remap=rm, **kwargs)
+        return ManifestLoader.Config(module_remap=rm)
 
-    #
+    ##
+
+    ENTRY_POINT_GROUP: ta.ClassVar[str] = 'omlish.manifests'
+
+    _discovered_packages: ta.ClassVar[ta.Optional[ta.Sequence[str]]] = None
+
+    @classmethod
+    def _discover_packages_uncached(cls) -> ta.Sequence[str]:
+        from importlib import metadata as importlib_metadata  # noqa
+        return [
+            ep.value
+            for ep in importlib_metadata.entry_points(group=cls.ENTRY_POINT_GROUP)
+        ]
+
+    @classmethod
+    def discover_packages(cls) -> ta.Sequence[str]:
+        if (x := cls._discovered_packages) is not None:
+            return x
+
+        x = cls._discover_packages_uncached()
+        cls._discovered_packages = x
+        return x
+
+    ##
+
+    def _scan_package_root_dir_uncached(
+            self,
+            root_dir: str,
+    ) -> ta.Sequence[str]:
+        pkgs: ta.List[str] = []
+
+        for n in os.listdir(root_dir):
+            if (
+                    os.path.isdir(p := os.path.join(root_dir, n)) and
+                    os.path.exists(os.path.join(p, '__init__.py'))
+            ):
+                pkgs.append(n)
+
+        return pkgs
+
+    def _scan_package_root_dir_locked(
+            self,
+            root_dir: str,
+    ) -> ta.Sequence[str]:
+        try:
+            return self._scanned_package_root_dirs[root_dir]
+        except KeyError:
+            pass
+
+        ret = self._scan_package_root_dir_uncached(root_dir)
+        self._scanned_package_root_dirs[root_dir] = ret
+        return ret
+
+    def _scan_package_root_dir(
+            self,
+            root_dir: str,
+    ) -> ta.Sequence[str]:
+        with self._lock:
+            return self._scan_package_root_dir_locked(root_dir)
+
+    ##
+
+    _detected_packages: ta.Set[str]
+
+    def _do_initialize(self) -> None:
+        self._detected_packages = set()
+
+        for r in self._config.package_scan_root_dirs or []:
+            self._detected_packages.update(self._scan_package_root_dir_locked(r))
+
+        if self._config.discover_packages:
+            self._detected_packages.update(dps := self.discover_packages())
+            if not dps:
+                for r in self._config.discover_packages_fallback_scan_root_dirs or []:
+                    self._detected_packages.update(self._scan_package_root_dir_locked(r))
+
+    _has_initialized = False
+
+    def _initialize_locked(self) -> None:
+        if not self._has_initialized:
+            self._do_initialize()
+            self._has_initialized = True
+
+    def has_initialized(self) -> bool:
+        with self._lock:
+            return self._has_initialized
+
+    def initialize(self) -> None:
+        if not self._has_initialized:
+            with self._lock:
+                self._initialize_locked()
+
+    ##
+
+    @dc.dataclass()
+    class ClassKeyError(Exception):
+        key: str
+        module: ta.Optional[str] = None
 
     def _load_class_uncached(self, key: str) -> type:
-        if not key.startswith('$'):
-            raise Exception(f'Bad key: {key}')
+        if not key.startswith('!'):
+            raise ManifestLoader.ClassKeyError(key)
 
         parts = key[1:].split('.')
+        if '' in parts:
+            raise ManifestLoader.ClassKeyError(key)
+
         pos = next(i for i, p in enumerate(parts) if p[0].isupper())
 
         mod_name = '.'.join(parts[:pos])
@@ -172,35 +337,84 @@ class ManifestLoader:
         with self._lock:
             return self._load_class_locked(key)
 
-    #
+    def get_class_key(self, cls: type) -> str:
+        if not (isinstance(cls, type) and dc.is_dataclass(cls)):
+            raise TypeError(cls)
+        mod_name = cls.__module__
+        mod_name = self._module_reverse_remap.get(mod_name, mod_name)
+        return f'!{mod_name}.{cls.__qualname__}'
+
+    ##
+
+    class _ResolvedManifest(ta.NamedTuple):
+        manifest: Manifest
+        package: str
+
+        module: str
+        class_key: str
+        value_dct: ta.Any
 
-    def _deserialize_raw_manifests(self, obj: ta.Any, pkg_name: str) -> ta.Sequence[Manifest]:
+    @classmethod
+    def _resolve_raw_manifest(
+            cls,
+            m: Manifest,
+            *,
+            package_name: str,
+    ) -> _ResolvedManifest:
+        # self._module = module
+        # self._class_key = class_key
+        if not m.module.startswith('.'):
+            raise NameError(m.module)
+
+        module = package_name + m.module
+
+        [(class_key, value_dct)] = m.value.items()
+
+        if not class_key.startswith('!'):
+            raise ManifestLoader.ClassKeyError(class_key, module=module)
+
+        if class_key.startswith('!.'):
+            class_key = f'!{package_name}{class_key[1:]}'
+
+        # FIXME: move to builder
+        # elif key.startswith('$.'):
+        #     if module.startswith('.'):
+        #         raise NameError(module)
+        #     kl = key[1:].split('.')
+        #     ml = module.split('.')
+        #     lvl = len(kl) - kl[::-1].index('')
+        #     if lvl >= len(ml):
+        #         raise ManifestLoader.ClassKeyError(key, module=module)
+        #     rn = '.'.join([*ml[:-lvl], *kl[lvl:]])
+        #     return f'${rn}'
+
+        return ManifestLoader._ResolvedManifest(
+            manifest=m,
+            package=package_name,
+
+            module=module,
+            class_key=class_key,
+            value_dct=value_dct,
+        )
+
+    ##
+
+    @classmethod
+    def _deserialize_raw_manifests(cls, obj: ta.Any) -> ta.Sequence[Manifest]:
         if not isinstance(obj, (list, tuple)):
             raise TypeError(obj)
 
         lst: ta.List[Manifest] = []
         for e in obj:
-            m = Manifest(**e)
-
-            m = dc.replace(m, module=pkg_name + m.module)
-
-            [(key, value_dct)] = m.value.items()
-            if not key.startswith('$'):
-                raise Exception(f'Bad key: {key}')
-            if key.startswith('$.'):
-                key = f'${pkg_name}{key[1:]}'
-                m = dc.replace(m, value={key: value_dct})
-
-            lst.append(m)
+            lst.append(Manifest(**e))
 
         return lst
 
-    #
-
-    def _read_package_file_text(self, pkg_name: str, file_name: str) -> ta.Optional[str]:
+    @classmethod
+    def _read_package_file_text(cls, package_name: str, file_name: str) -> ta.Optional[str]:
         # importlib.resources.files actually imports the package - to avoid this, if possible, the file is read straight
         # off the filesystem.
-        spec = importlib.util.find_spec(pkg_name)
+        spec = importlib.util.find_spec(package_name)
         if (
                 spec is not None and
                 isinstance(spec.loader, importlib.machinery.SourceFileLoader) and
@@ -214,15 +428,17 @@ class ManifestLoader:
                 with open(file_path, encoding='utf-8') as f:
                     return f.read()
 
-        t = importlib.resources.files(pkg_name).joinpath(file_name)
+        from importlib import resources as importlib_resources  # noqa
+        t = importlib_resources.files(package_name).joinpath(file_name)
         if not t.is_file():
             return None
         return t.read_text('utf-8')
 
     MANIFESTS_FILE_NAME: ta.ClassVar[str] = '.manifests.json'
 
-    def _load_package_uncached(self, pkg_name: str) -> ta.Optional[LoadedPackage]:
-        src = self._read_package_file_text(pkg_name, self.MANIFESTS_FILE_NAME)
+    @classmethod
+    def _read_package_raw_manifests(cls, package_name: str) -> ta.Optional[ta.Sequence[Manifest]]:
+        src = cls._read_package_file_text(package_name, cls.MANIFESTS_FILE_NAME)
         if src is None:
             return None
 
@@ -230,13 +446,21 @@ class ManifestLoader:
         if not isinstance(obj, (list, tuple)):
             raise TypeError(obj)
 
-        raw_lst = self._deserialize_raw_manifests(obj, pkg_name)
+        return cls._deserialize_raw_manifests(obj)
 
-        ld_pkg = ManifestLoader.LoadedPackage(self, pkg_name)
+    ##
+
+    def _load_package_uncached(self, package_name: str) -> ta.Optional[LoadedPackage]:
+        if (raw_lst := self._read_package_raw_manifests(package_name)) is None:
+            return None
+
+        ld_pkg = ManifestLoader.LoadedPackage(self, package_name)
 
         ld_man_lst: ta.List[ManifestLoader.LoadedManifest] = []
         for raw in raw_lst:
-            ld_man = ManifestLoader.LoadedManifest(ld_pkg, raw)
+            rsv = self._resolve_raw_manifest(raw, package_name=package_name)
+
+            ld_man = ManifestLoader.LoadedManifest(ld_pkg, rsv)
 
             ld_man_lst.append(ld_man)
 
@@ -244,167 +468,116 @@ class ManifestLoader:
 
         return ld_pkg
 
-    def _load_package_locked(self, pkg_name: str) -> ta.Optional[LoadedPackage]:
+    def _load_package_locked(self, package_name: str) -> ta.Optional[LoadedPackage]:
         try:
-            return self._loaded_packages[pkg_name]
+            return self._loaded_packages[package_name]
         except KeyError:
             pass
 
-        pkg = self._load_package_uncached(pkg_name)
-        self._loaded_packages[pkg_name] = pkg
+        pkg = self._load_package_uncached(package_name)
+        self._loaded_packages[package_name] = pkg
         return pkg
 
-    def load_package(self, pkg_name: str) -> ta.Optional[LoadedPackage]:
+    def load_package(self, package_name: str) -> ta.Optional[LoadedPackage]:
         with self._lock:
-            return self._load_package_locked(pkg_name)
+            return self._load_package_locked(package_name)
 
-    #
+    ##
 
     def _instantiate_value(self, cls: type, **kwargs: ta.Any) -> ta.Any:
-        if self._value_instantiator is not None:
-            return self._value_instantiator(cls, **kwargs)
+        if self._config.value_instantiator is not None:
+            return self._config.value_instantiator(cls, **kwargs)
         else:
             return cls(**kwargs)
 
-    def _instantiate_loaded_manifest(self, ld_man: LoadedManifest) -> ta.Any:
-        [(cls_key, value_dct)] = ld_man.manifest.value.items()
-        cls = self._load_class(cls_key)
-        value = self._instantiate_value(cls, **value_dct)
+    def _instantiate_resolved_manifest(self, resolved: _ResolvedManifest) -> ta.Any:
+        cls = self._load_class(resolved.class_key)
+        value = self._instantiate_value(cls, **resolved.value_dct)
         return value
 
-    #
-
-    # FIXME:
-    # class LOAD_ALL:  # noqa
-    #     def __new__(cls, *args, **kwargs):  # noqa
-    #         raise TypeError
-    #
-    #     def __init_subclass__(cls, **kwargs):  # noqa
-    #         raise TypeError
+    ##
 
-    def _load(
+    def _load_initialized(
             self,
-            *pkg_names: str,
-            only: ta.Optional[ta.Iterable[type]] = None,
+            *,
+            packages: ta.Optional[ta.Collection[str]] = None,
+            classes: ta.Optional[ta.Collection[type]] = None,
     ) -> ta.Sequence[LoadedManifest]:
-        only_keys: ta.Optional[ta.Set]
-        if only is not None:
-            only_keys = set()
-            for cls in only:
-                if not (isinstance(cls, type) and dc.is_dataclass(cls)):
-                    raise TypeError(cls)
-                mod_name = cls.__module__
-                mod_name = self._module_reverse_remap.get(mod_name, mod_name)
-                only_keys.add(f'${mod_name}.{cls.__qualname__}')
-        else:
-            only_keys = None
+        if isinstance(packages, str):
+            raise TypeError(packages)
+
+        class_keys: ta.Optional[ta.Set[str]] = None
+        if classes is not None:
+            class_keys = {self.get_class_key(cls) for cls in classes}
+
+        if packages is None:
+            packages = self._detected_packages
 
         lst: ta.List[ManifestLoader.LoadedManifest] = []
-        for pn in pkg_names:
-            lp = self.load_package(pn)
-            if lp is None:
+        for pn in packages:
+            if (lp := self._load_package_locked(pn)) is None:
                 continue
 
-            for m in lp.manifests:
-                if only_keys is not None and m.class_key not in only_keys:
-                    continue
+            if class_keys is not None:
+                for ck in class_keys:
+                    lst.extend(lp.manifests_by_class_key.get(ck, []))
 
-                lst.append(m)
+            else:
+                lst.extend(lp.manifests)
 
         return lst
 
+    def _load_locked(
+            self,
+            *,
+            packages: ta.Optional[ta.Collection[str]] = None,
+            classes: ta.Optional[ta.Collection[type]] = None,
+    ) -> ta.Sequence[LoadedManifest]:
+        self._initialize_locked()
+        return self._load_initialized(
+            packages=packages,
+            classes=classes,
+        )
+
     def load(
             self,
-            *pkg_names: str,
-            only: ta.Optional[ta.Iterable[type]] = None,
+            *,
+            packages: ta.Optional[ta.Collection[str]] = None,
+            classes: ta.Optional[ta.Collection[type]] = None,
     ) -> ta.Sequence[LoadedManifest]:
+        if isinstance(packages, str):
+            raise TypeError(packages)
+
         with self._lock:
-            return self._load(
-                *pkg_names,
-                only=only,
+            return self._load_locked(
+                packages=packages,
+                classes=classes,
             )
 
-    ##
-
-    ENTRY_POINT_GROUP: ta.ClassVar[str] = 'omlish.manifests'
-
-    _discovered_packages: ta.ClassVar[ta.Optional[ta.Sequence[str]]] = None
-
-    @classmethod
-    def _discover_packages_uncached(cls) -> ta.Sequence[str]:
-        # This is a fat dep so do it late.
-        from importlib import metadata as importlib_metadata  # noqa
-
+    def load_values(
+            self,
+            *,
+            packages: ta.Optional[ta.Collection[str]] = None,
+            classes: ta.Optional[ta.Collection[type]] = None,
+    ) -> ta.Sequence[ta.Any]:
         return [
-            ep.value
-            for ep in importlib_metadata.entry_points(group=cls.ENTRY_POINT_GROUP)
+            lm.value()
+            for lm in self.load(
+                packages=packages,
+                classes=classes,
+            )
         ]
 
-    @classmethod
-    def discover_packages(cls) -> ta.Sequence[str]:
-        if (x := cls._discovered_packages) is not None:
-            return x
-
-        x = cls._discover_packages_uncached()
-        cls._discovered_packages = x
-        return x
-
-    ##
-
-    def _scan_package_root_dir_uncached(
-            self,
-            root_dir: str,
-    ) -> ta.Sequence[str]:
-        pkgs: ta.List[str] = []
-
-        for n in os.listdir(root_dir):
-            if (
-                    os.path.isdir(p := os.path.join(root_dir, n)) and
-                    os.path.exists(os.path.join(p, '__init__.py'))
-            ):
-                pkgs.append(n)
-
-        return pkgs
-
-    def _scan_package_root_dir_locked(
-            self,
-            root_dir: str,
-    ) -> ta.Sequence[str]:
-        try:
-            return self._scanned_package_root_dirs[root_dir]
-        except KeyError:
-            pass
-
-        ret = self._scan_package_root_dir_uncached(root_dir)
-        self._scanned_package_root_dirs[root_dir] = ret
-        return ret
-
-    def _scan_package_root_dir(
-            self,
-            root_dir: str,
-    ) -> ta.Sequence[str]:
-        with self._lock:
-            return self._scan_package_root_dir_locked(root_dir)
-
-    def scan_or_discover_packages(
+    def load_values_of(
             self,
+            cls: ta.Type[T],
             *,
-            specified_root_dirs: ta.Optional[ta.Sequence[str]] = None,
-            fallback_root_dir: ta.Optional[str] = None,
-    ) -> ta.Sequence[str]:
-        pkgs: list[str] = []
-
-        if specified_root_dirs is not None:
-            if isinstance(specified_root_dirs, str):
-                raise TypeError(specified_root_dirs)
-
-            for r in specified_root_dirs:
-                pkgs.extend(self._scan_package_root_dir(r))
-
-        else:
-            pkgs.extend(self.discover_packages())
-
-            if not pkgs and fallback_root_dir is not None:
-                pkgs.extend(self._scan_package_root_dir(fallback_root_dir))
-
-        return pkgs
+            packages: ta.Optional[ta.Collection[str]] = None,
+    ) -> ta.Sequence[T]:
+        return [
+            ta.cast(T, lm.value())
+            for lm in self.load(
+                packages=packages,
+                classes=[cls],
+            )
+        ]