Migrating the classes to pydantic

playbook/action_registry/__init__.py

@@ -1,49 +1,58 @@
-from dataclasses import dataclass
-from playbook.models import StepLog, ActionFn
-from typing import Callable, Dict, Any
+import os
+import sys
+import importlib.util
+from playbook.logging_models import StepLogModel
+from typing import Callable
+from pydantic import BaseModel, PrivateAttr
 
 
-@dataclass
-class Action(object):
+ActionFn = Callable[[dict[str, object], str], StepLogModel]
+
+
+class Action(BaseModel):
     name: str
     fn: ActionFn
     ver: str
 
 
-class ActionRegistry:
-    def __init__(self, name: str):
-        self.name: str = name
-        self.__actions = {}
+class ActionRegistry(BaseModel):
+    name: str
+    _actions: dict[str, Action] = PrivateAttr(default_factory=dict)
 
     def register(self, name: str, version: str):
         """Decorator to register python functions with a name and version."""
         def decorator(function: ActionFn):
-            if name not in self.__actions:
-                self.__actions[name] = {}
-
-            self.__actions[name][version] = function
+            self._actions[name] = Action(name=name, fn=function, ver=version)
             return function
-
         return decorator
-    
-    def get(self, name: str, ver: str) -> ActionFn:
-        if name not in self.__actions:
-            raise ValueError(f"Action '{name}' is not registered in registry '{self.name}'")
-        
-        versions_dict = self.__actions[name]
-        if not versions_dict:
-            raise ValueError(f"No registered versions found for function '{name}'")
 
-        if ver == "*":
-            return next(iter(versions_dict.values()))
+    def get(self, name: str, ver: str = '*') -> ActionFn:
+        if name not in self._actions:
+            raise ValueError(
+                f"Action '{name}' is not registered in registry '{self.name}'")
 
-        if ver not in self.__actions[name]:
-            raise ValueError(f"Version '{ver}' is not registered for function '{name}'")
+        return self._actions[name].fn
 
-        return versions_dict[ver]
+    @staticmethod
+    def load_registries_from_file(file_path: str) -> list[ActionRegistry]:
+        abs_path = os.path.abspath(file_path)
+        module_name = os.path.splitext(os.path.basename(abs_path))[0]
 
-    def manifest(self):
-        functions = {}
-        for fn_name, versions_dict in self.__actions.items():
-            functions[fn_name] = list(versions_dict.keys())
-        return {"registry_name": self.name, "functions": functions}
+        spec = importlib.util.spec_from_file_location(module_name, abs_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(
+                f"could not create spec for registry file: {file_path}")
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+
+        discovered_registries = [
+            obj for obj in vars(module).values()
+            if isinstance(obj, ActionRegistry)
+        ]
+
+        if not discovered_registries:
+            raise ValueError(f"No ActionRegistry instances were found in '{file_path}'")
+
+        return discovered_registries