import inspect import json import os import pkgutil import re from abc import ABC, abstractmethod from dataclasses import asdict, dataclass, field, fields from typing import final class AbstractField: pass from .text_utils import camel_to_snake_case, is_camel_case artifactories = [] class BaseArtifact(ABC): _class_register = {} @classmethod def is_artifact_dict(cls, d): return isinstance(d, dict) and "type" in d and d["type"] in cls._class_register @classmethod def register_class(cls, artifact_class): assert issubclass(artifact_class, BaseArtifact), "Artifact class must be a subclass of BaseArtifact" assert is_camel_case( artifact_class.__name__ ), f"Artifact class name must be legal camel case, got {artifact_class.__name__}" snake_case_key = camel_to_snake_case(artifact_class.__name__) if snake_case_key in cls._class_register: assert ( cls._class_register[snake_case_key] == artifact_class ), f"Artifact class name must be unique, {snake_case_key} already exists for {cls._class_register[snake_case_key]}" cls._class_register[snake_case_key] = artifact_class return snake_case_key @classmethod def is_artifact_file(cls, path): if not os.path.exists(path) or not os.path.isfile(path): return False with open(path, "r") as f: d = json.load(f) return cls.is_artifact_dict(d) @final def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @final def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) cls = dataclass(cls) def prepare(self): pass def verify(self): pass @final def __post_init__(self): self.type = self.register_class(self.__class__) self._args_dict = asdict(self) for field in fields(self): # check if field.type is class and if it is subclass of BaseArtifact if isinstance(field.type, type) and issubclass(field.type, BaseArtifact): value = getattr(self, field.name) if isinstance(value, str): artifact, artifactory = fetch_artifact(value) assert artifact is not None, f"Artifact {value} does not exist, in {artifactories}" print(f"Artifact {value} is fetched from {artifactory}") setattr(self, field.name, artifact) self.prepare() self.verify() def to_dict(self): return self._args_dict def save(self, path): with open(path, "w") as f: json.dump(self.to_dict(), f, indent=4) # def __getstate__(self): # print('getstate', self.__dict__) # return self.to_dict() @classmethod def _recursive_load(cls, d): if isinstance(d, dict): new_d = {} for key, value in d.items(): new_d[key] = cls._recursive_load(value) d = new_d elif isinstance(d, list): d = [cls._recursive_load(value) for value in d] else: pass if cls.is_artifact_dict(d): instance = cls._class_register[d.pop("type")](**d) return instance else: return d @classmethod def from_dict(cls, d): assert cls.is_artifact_dict(d), "Input must be a dict with type field" return cls._recursive_load(d) @classmethod def load(cls, path): with open(path, "r") as f: d = json.load(f) assert "type" in d, "Saved artifact must have a type field" return cls._recursive_load(d) # assert d['type'] in cls._class_register, f'Artifact type "{d["type"]}" is not registered' # cls = cls._class_register[d.pop('type')] # return cls(**d) class Artifact(BaseArtifact): type: str = field(init=False) class ArtifactList(list, Artifact): def prepare(self): for artifact in self: artifact.prepare() class Artifactory(Artifact, ABC): @abstractmethod def __contains__(self, name: str) -> bool: pass @abstractmethod def __getitem__(self, name) -> Artifact: pass class UnitxtArtifactNotFoundError(Exception): def __init__(self, name, artifactories): self.name = name self.artifactories = artifactories def __str__(self): return f"Artifact {self.name} does not exist, in artifactories:{self.artifactories}" def fetch_artifact(name): if Artifact.is_artifact_file(name): return Artifact.load(name), None else: for artifactory in artifactories: if name in artifactory: return artifactory[name], artifactory raise UnitxtArtifactNotFoundError(name, artifactories) def register_atrifactory(artifactory): assert isinstance(artifactory, Artifactory), "Artifactory must be an instance of Artifactory" assert hasattr(artifactory, "__contains__"), "Artifactory must have __contains__ method" assert hasattr(artifactory, "__getitem__"), "Artifactory must have __getitem__ method" artifactories.append(artifactory) def register_all_artifacts(path): for loader, module_name, is_pkg in pkgutil.walk_packages(path): print(__name__) if module_name == __name__: continue print(f"Loading {module_name}") # Import the module module = loader.find_module(module_name).load_module(module_name) # Iterate over every object in the module for name, obj in inspect.getmembers(module): # Make sure the object is a class if inspect.isclass(obj): # Make sure the class is a subclass of Artifact (but not Artifact itself) if issubclass(obj, Artifact) and obj is not Artifact: print(obj)