File size: 5,949 Bytes
c6d1c21 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
from dataclasses import dataclass, asdict, fields, field
from abc import ABC, abstractmethod
from typing import final
import re
import json
import os
import pkgutil
import inspect
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)
|