Resolve #54: Exposes resource_registry.register(…) to register user…

… defined custom resources.

docs/codecs.md

@@ -41,7 +41,8 @@ Output: `<class 'lightkube.resources.core_v1.ConfigMap'>`
 
 !!! note
     Only known resources can be loaded. These are either kubernetes [standard resources](resources-and-models.md) 
-    or [generic resources](generic-resources.md) manually defined.
+    or [generic resources](generic-resources.md) manually defined. You can register further resources using
+    the [`resource_registry`](#resource-registry).
 
 ## Load from YAML
 
@@ -212,3 +213,34 @@ with open('crs_amd_crds.yaml') as f:
 ```
 
 This orders the objects in a way that is friendly for deleting them as a batch.
+
+## Resource Registry
+
+The singleton `resource_registry` allows to register a custom resource, so that it can be used by the load
+functions on this module:
+
+```python
+from lightkube import codecs
+
+codecs.resource_registry.register(MyCustomResource)
+
+with open('service.yaml') as f:
+    # Now `MyCustomResource` can be loaded
+    objs = codecs.load_all_yaml(f)
+```
+
+`register` can also be used as a decorator:
+```python
+from lightkube.core.resource import NamespacedResource
+from lightkube.codecs import resource_registry
+
+@resource_registry.register
+class MyCustomResource(NamespacedResource):
+    ...
+```
+
+### Reference
+
+::: lightkube.codecs.resource_registry
+    :docstring:
+    :members:

lightkube/codecs.py

@@ -1,11 +1,17 @@
-import importlib
 from typing import Union, TextIO, Iterator, List, Mapping
 
 import yaml
 
-from .generic_resource import get_generic_resource, GenericGlobalResource, GenericNamespacedResource, create_resources_from_crd
+from .generic_resource import GenericGlobalResource, GenericNamespacedResource, create_resources_from_crd
 from .core.exceptions import LoadResourceError
+from .core.resource_registry import resource_registry
 
+__all__ = [
+    'from_dict',
+    'load_all_yaml',
+    'dump_all_yaml',
+    'resource_registry'
+]
 
 try:
     import jinja2
@@ -17,30 +23,6 @@
 AnyResource = Union[GenericGlobalResource, GenericNamespacedResource]
 
 
-def _load_model(version, kind):
-    if "/" in version:
-        group, version_n = version.split("/")
-        # Check if a generic resource was defined
-        model = get_generic_resource(version, kind)
-        if model is not None:
-            return model
-
-        # Generic resource not defined, but it could be a k8s resource
-        if group.endswith(".k8s.io"):
-            group = group[:-7]
-        group = group.replace(".", "_")
-        version = "_".join([group, version_n])
-    else:
-        version = f'core_{version}'
-
-    try:
-        module = importlib.import_module(f'lightkube.resources.{version.lower()}')
-    except ImportError as e:
-        # It was not a k8s resource and a generic resource was not previously defined
-        raise LoadResourceError(f"{e}. If using a CRD, ensure you define a generic resource.")
-    return getattr(module, kind)
-
-
 def from_dict(d: dict) -> AnyResource:
     """Converts a kubernetes resource defined as python dict to the corresponding resource object.
     If the dict represent a standard resource, the function will automatically load the appropriate
@@ -58,7 +40,7 @@ def from_dict(d: dict) -> AnyResource:
         if attr not in d:
             raise LoadResourceError(f"Invalid resource definition, key '{attr}' missing.")
 
-    model = _load_model(d['apiVersion'], d['kind'])
+    model = resource_registry.load(d['apiVersion'], d['kind'])
     return model.from_dict(d)
 
 

lightkube/core/resource_registry.py

@@ -0,0 +1,106 @@
+import importlib
+from typing import Union, Type, Optional
+
+from lightkube.core import resource as res
+from lightkube.core.exceptions import LoadResourceError
+
+AnyResource = Union[res.NamespacedResource, res.GlobalResource]
+
+def _load_internal_resource(version, kind):
+    if "/" in version:
+        group, version_n = version.split("/")
+        # Generic resource not defined, but it could be a k8s resource
+        if group.endswith(".k8s.io"):
+            group = group[:-7]
+        group = group.replace(".", "_")
+        module_name = "_".join([group, version_n])
+    else:
+        module_name = f'core_{version}'
+
+    module = importlib.import_module(f'lightkube.resources.{module_name.lower()}')
+    try:
+        return getattr(module, kind)
+    except AttributeError:
+        raise LoadResourceError(f"Cannot find resource kind '{kind}' in module {module.__name__}")
+
+def _maybe_internal(version):
+    if "/" not in version:
+        return True
+
+    group = version.split("/")[0]
+    # internal resources don't have namespace or end in .k8s.io
+    return group.endswith(".k8s.io") or "." not in group
+
+
+class ResourceRegistry:
+    """Resource Registry used to load standard resources or to register custom resources
+    """
+    _registry: dict
+
+    def __init__(self):
+        self._registry = {}
+
+    def register(self, resource: Type[AnyResource]) -> Type[AnyResource]:
+        """Register a custom resource
+
+        **parameters**
+
+        * **resource** - Resource class to register.
+
+        **returns** The `resource` class provided
+        """
+        info = resource._api_info
+        version = f'{info.resource.group}/{info.resource.version}' if info.resource.group else info.resource.version
+        res_key = (version, info.resource.kind)
+
+        if res_key in self._registry:
+            registered_resource = self._registry[res_key]
+            if registered_resource is resource:   # already present
+                return registered_resource
+            raise ValueError(f"Another class for resource '{info.resource.kind}' is already registered")
+
+        self._registry[res_key] = resource
+        return resource
+
+    def clear(self):
+        """Clear the registry from all registered resources
+        """
+        self._registry.clear()
+
+    def get(self, version: str, kind: str) -> Optional[Type[AnyResource]]:
+        """Get a resource from the registry matching the given `version` and `kind`.
+
+        **parameters**
+
+        * **version** - Version of the resource as defined in the kubernetes definition. Example `example.com/v1`
+        * **kind** - Resource kind. Example `CronJob`
+
+        **returns** A `resource` class or `None` if there is no match in the registry.
+        """
+        return self._registry.get((version, kind))
+
+    def load(self, version, kind) -> Optional[Type[AnyResource]]:
+        """Load a standard resource from `lightkube.resources` given `version` and `kind`.
+        This method look up the registry first and import the resource from the module only if it's not available there.
+
+        * **version** - Version of the resource as defined in the kubernetes definition. Example `apps/v1`
+        * **kind** - Resource kind. Example `Pod`
+
+        **returns** A `resource` class if the resource is found. Otherwise an exception is raised
+        """
+        # check if this resource is in the registry
+        resource = self.get(version, kind)
+        if resource is not None:
+            return resource
+
+        # if not, we attempt to load from lightkube.resources
+        if _maybe_internal(version):
+            try:
+                return self.register(_load_internal_resource(version, kind))
+            except ImportError:
+                pass
+
+        raise LoadResourceError(f"Cannot find resource {kind} of group {version}. "
+                                "If using a CRD, ensure a generic resource is defined.")
+
+resource_registry = ResourceRegistry()

lightkube/generic_resource.py

@@ -5,6 +5,7 @@
 from .core.async_client import AsyncClient
 from .core.internal_models import meta_v1, autoscaling_v1
 from .core.internal_resources import apiextensions
+from .core.resource_registry import resource_registry
 
 
 __all__ = [
@@ -16,8 +17,6 @@
     'load_in_cluster_generic_resources',
 ]
 
-_created_resources = {}
-
 
 def get_generic_resource(version, kind):
     """Query generic resources already defined using one of the other methods described in this module or via
@@ -30,9 +29,10 @@ def get_generic_resource(version, kind):
 
     **returns** class representing the generic resource or `None` if it's not found
     """
-    global _created_resources
-    model = _created_resources.get((version, kind))
-    return model[0] if model is not None else None
+    resource = resource_registry.get(version, kind)
+    if resource is None or not issubclass(resource, (GenericGlobalResource, GenericNamespacedResource)):
+        return None
+    return resource
 
 
 class Generic(dict):
@@ -132,12 +132,18 @@ class GenericNamespacedResource(res.NamespacedResourceG, Generic):
     Status: Type[GenericNamespacedStatus]
 
 
+def _api_info_signature(api_info: res.ApiInfo, namespaced: bool):
+    return (namespaced, api_info.plural, tuple(api_info.verbs) if api_info.verbs else None)
+
+
 def _create_resource(namespaced, group, version, kind, plural, verbs=None) -> Any:
-    global _created_resources
-    res_key = (f'{group}/{version}', kind)
-    signature = (namespaced, plural, tuple(verbs) if verbs else None)
-    if res_key in _created_resources:
-        model, curr_signature = _created_resources[res_key]
+    model = resource_registry.get(f'{group}/{version}', kind)
+    api_info = create_api_info(group, version, kind, plural, verbs=verbs)
+    signature = _api_info_signature(api_info, namespaced)
+
+    if model is not None:
+        curr_namespaced = issubclass(model, res.NamespacedResource)
+        curr_signature = _api_info_signature(model._api_info, curr_namespaced)
         if curr_signature != signature:
             raise ValueError(f"Resource {kind} already created but with different signature")
         return model
@@ -154,8 +160,7 @@ class TmpName(main):
         Status = _create_subresource(status, _api_info, action='status')
 
     TmpName.__name__ = TmpName.__qualname__ = kind
-    _created_resources[res_key] = (TmpName, signature)
-    return TmpName
+    return resource_registry.register(TmpName)
 
 
 def create_global_resource(group: str, version: str, kind: str, plural: str, verbs=None) \

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name='lightkube',
-    version="0.14.0",
+    version="0.15.0",
     description='Lightweight kubernetes client library',
     long_description=Path("README.md").read_text(),
     long_description_content_type="text/markdown",

tests/test_codecs.py

@@ -11,16 +11,16 @@
 from lightkube.models.meta_v1 import ObjectMeta
 from lightkube import generic_resource as gr
 from lightkube import LoadResourceError
+from lightkube.codecs import resource_registry
 
 data_dir = Path(__file__).parent.joinpath('data')
 
 
-@pytest.fixture()
-def mocked_created_resources():
-    """Isolates the module variable generic_resource._created_resources to avoid test spillover"""
-    created_resources_dict = {}
-    with mock.patch("lightkube.generic_resource._created_resources", created_resources_dict) as mocked_created_resources:
-        yield mocked_created_resources
+@pytest.fixture(autouse=True)
+def cleanup_registry():
+    """Cleanup the registry before each test"""
+    yield
+    resource_registry.clear()
 
 
 def test_from_dict():
@@ -111,7 +111,7 @@ def test_from_dict_not_found():
     with pytest.raises(LoadResourceError):
         codecs.from_dict({'apiVersion': 'myapp2.com/v1', 'kind': 'Mydb'})
 
-    with pytest.raises(AttributeError):
+    with pytest.raises(LoadResourceError):
         codecs.from_dict({'apiVersion': 'v1', 'kind': 'Missing'})
 
     with pytest.raises(LoadResourceError):
@@ -131,6 +131,7 @@ def test_from_dict_not_found():
     )
 )
 def test_load_all_yaml_static(yaml_file):
+    gr.create_namespaced_resource('myapp.com', 'v1', 'Mydb', 'mydbs')
     objs = list(codecs.load_all_yaml(data_dir.joinpath(yaml_file).read_text()))
     kinds = [o.kind for o in objs]
 
@@ -144,6 +145,7 @@ def test_load_all_yaml_static(yaml_file):
 
 
 def test_load_all_yaml_template():
+    gr.create_namespaced_resource('myapp.com', 'v1', 'Mydb', 'mydbs')
     objs = list(codecs.load_all_yaml(
         data_dir.joinpath('example-def.tmpl').read_text(),
         context={'test': 'xyz'})
@@ -162,6 +164,7 @@ def test_load_all_yaml_template():
 
 
 def test_load_all_yaml_template_env():
+    gr.create_namespaced_resource('myapp.com', 'v1', 'Mydb', 'mydbs')
     import jinja2
     env = jinja2.Environment()
     env.globals['test'] = 'global'
@@ -214,15 +217,15 @@ def test_load_all_yaml_missing_dependency():
         False,  # no generic resources should be created
     ]
 )
-def test_load_all_yaml_creating_generic_resources(create_resources_for_crds, mocked_created_resources):
+def test_load_all_yaml_creating_generic_resources(create_resources_for_crds):
     template_yaml = data_dir.joinpath('example-multi-version-crd.yaml').read_text()
     template_dict = list(yaml.safe_load_all(template_yaml))[0]
 
     expected_group = template_dict["spec"]["group"]
     expected_kind = template_dict["spec"]["names"]["kind"]
 
     # Confirm no generic resources exist before testing
-    assert len(gr._created_resources) == 0
+    assert len(resource_registry._registry) == 0
 
     objs = list(codecs.load_all_yaml(
         template_yaml,
@@ -232,14 +235,15 @@ def test_load_all_yaml_creating_generic_resources(create_resources_for_crds, moc
     # Confirm expected resources exist
     if create_resources_for_crds:
         for version in template_dict["spec"]["versions"]:
-            resource = gr.get_generic_resource(f"{expected_group}/{version['name']}", expected_kind)
+            resource = resource_registry.get(f"{expected_group}/{version['name']}", expected_kind)
             assert resource is not None
 
         # Confirm we did not make any extra resources
-        assert len(gr._created_resources) == len(template_dict["spec"]["versions"])
+        # + 1 as CustomResourceDefinition is also added to the registry
+        assert len(resource_registry._registry) == len(template_dict["spec"]["versions"]) + 1
     else:
-        # Confirm we did not make any resources
-        assert len(gr._created_resources) == 0
+        # Confirm we did not make any resources except CustomResourceDefinition
+        assert len(resource_registry._registry) == 1
 
     assert len(objs) == 1