File size: 9,851 Bytes
0e3c8e6
4d7f032
c6d1c21
 
 
e80560b
eac4eaf
f0b2749
e80560b
c6d1c21
4a664e8
ba3eb02
eac4eaf
 
ba3eb02
 
 
c6d1c21
 
e80560b
ba0637a
eac4eaf
e80560b
ba0637a
 
 
c6d1c21
ba0637a
 
 
 
 
 
214d47a
e80560b
 
 
 
 
 
 
 
 
 
c6d1c21
214d47a
e80560b
 
 
 
 
 
 
 
 
214d47a
 
 
 
 
c6d1c21
eac4eaf
 
 
 
 
 
 
 
 
 
 
 
0e3c8e6
 
 
 
 
 
 
 
e80560b
0e3c8e6
214d47a
 
0e3c8e6
214d47a
 
0e3c8e6
 
 
e80560b
0e3c8e6
e80560b
 
 
0e3c8e6
 
 
eac4eaf
 
 
c6d1c21
 
4a664e8
 
c6d1c21
 
214d47a
c6d1c21
0e3c8e6
214d47a
0e3c8e6
e80560b
 
 
0e3c8e6
e80560b
214d47a
e80560b
0e3c8e6
c3f3d16
 
 
 
c6d1c21
 
eac4eaf
 
ba3eb02
c6d1c21
 
ba3eb02
c6d1c21
 
 
214d47a
c6d1c21
7391dc4
 
c6d1c21
c3f3d16
 
c6d1c21
 
 
 
214d47a
 
 
 
c6d1c21
 
 
 
e80560b
c6d1c21
 
 
214d47a
 
 
 
ba3eb02
 
 
 
 
214d47a
 
 
 
c6d1c21
 
 
 
 
 
 
 
 
 
 
 
214d47a
e80560b
 
 
c6d1c21
 
 
214d47a
c6d1c21
 
 
4a664e8
ba3eb02
4a664e8
 
 
c6d1c21
eac4eaf
 
c6d1c21
eac4eaf
 
 
c3f3d16
 
287304a
c3f3d16
eac4eaf
 
 
 
 
e80560b
 
 
eac4eaf
 
 
 
 
 
 
c3f3d16
 
eac4eaf
 
ba3eb02
 
c6d1c21
 
287304a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c6d1c21
 
 
 
 
 
eac4eaf
c6d1c21
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f0b2749
c6d1c21
 
 
e80560b
 
 
 
c6d1c21
ba0637a
c6d1c21
eac4eaf
f0b2749
66c1161
eac4eaf
ba3eb02
eac4eaf
 
66c1161
f0b2749
 
 
 
 
66c1161
 
 
e80560b
 
c6d1c21
eac4eaf
c6d1c21
e80560b
ba3eb02
c6d1c21
 
ba3eb02
c6d1c21
 
 
 
e80560b
c6d1c21
 
 
 
ba3eb02
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
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
import difflib
import inspect
import json
import os
import pkgutil
from abc import abstractmethod
from copy import deepcopy
from functools import lru_cache
from typing import Dict, List, Union, final

from .dataclass import Dataclass, Field, InternalField, fields
from .logging_utils import get_logger
from .text_utils import camel_to_snake_case, is_camel_case
from .type_utils import issubtype
from .utils import load_json, save_json

logger = get_logger()


class Artifactories:
    def __new__(cls):
        if not hasattr(cls, "instance"):
            cls.instance = super().__new__(cls)
            cls.instance.artifactories = []

        return cls.instance

    def __iter__(self):
        return iter(self.artifactories)

    def __next__(self):
        return next(self.artifactories)

    def register(self, 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"
        self.artifactories = [artifactory, *self.artifactories]

    def unregister(self, 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"
        self.artifactories.remove(artifactory)

    def reset(self):
        self.artifactories = []


def map_values_in_place(object, mapper):
    if isinstance(object, dict):
        for key, value in object.items():
            object[key] = mapper(value)
        return object
    if isinstance(object, list):
        for i in range(len(object)):
            object[i] = mapper(object[i])
        return object
    return mapper(object)


def get_closest_artifact_type(type):
    artifact_type_options = list(Artifact._class_register.keys())
    matches = difflib.get_close_matches(type, artifact_type_options)
    if matches:
        return matches[0]  # Return the closest match
    return None


class UnrecognizedArtifactTypeError(ValueError):
    def __init__(self, type) -> None:
        maybe_class = "".join(word.capitalize() for word in type.split("_"))
        message = f"'{type}' is not a recognized artifact 'type'. Make sure a the class defined this type (Probably called '{maybe_class}' or similar) is defined and/or imported anywhere in the code executed."
        closest_artifact_type = get_closest_artifact_type(type)
        if closest_artifact_type is not None:
            message += "\n\n" f"Did you mean '{closest_artifact_type}'?"
        super().__init__(message)


class MissingArtifactTypeError(ValueError):
    def __init__(self, dic) -> None:
        message = (
            f"Missing 'type' parameter. Expected 'type' in artifact dict, got {dic}"
        )
        super().__init__(message)


class Artifact(Dataclass):
    type: str = Field(default=None, final=True, init=False)

    _class_register = {}

    artifact_identifier: str = InternalField(default=None, required=False)

    @classmethod
    def is_artifact_dict(cls, d):
        return isinstance(d, dict) and "type" in d

    @classmethod
    def verify_artifact_dict(cls, d):
        if not isinstance(d, dict):
            raise ValueError(
                f"Artifact dict <{d}> must be of type 'dict', got '{type(d)}'."
            )
        if "type" not in d:
            raise MissingArtifactTypeError(d)
        if not cls.is_registered_type(d["type"]):
            raise UnrecognizedArtifactTypeError(d["type"])

    @classmethod
    def get_artifact_type(cls):
        return camel_to_snake_case(cls.__name__)

    @classmethod
    def register_class(cls, artifact_class):
        assert issubclass(
            artifact_class, Artifact
        ), f"Artifact class must be a subclass of Artifact, got '{artifact_class}'"
        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 cls.is_registered_type(snake_case_key):
            assert (
                str(cls._class_register[snake_case_key]) == str(artifact_class)
            ), f"Artifact class name must be unique, '{snake_case_key}' already exists for {cls._class_register[snake_case_key]}. Cannot be overriden by {artifact_class}."

            return snake_case_key

        cls._class_register[snake_case_key] = artifact_class

        return snake_case_key

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls.register_class(cls)

    @classmethod
    def is_artifact_file(cls, path):
        if not os.path.exists(path) or not os.path.isfile(path):
            return False
        with open(path) as f:
            d = json.load(f)
        return cls.is_artifact_dict(d)

    @classmethod
    def is_registered_type(cls, type: str):
        return type in cls._class_register

    @classmethod
    def is_registered_class_name(cls, class_name: str):
        snake_case_key = camel_to_snake_case(class_name)
        return cls.is_registered_type(snake_case_key)

    @classmethod
    def is_registered_class(cls, clz: object):
        return clz in set(cls._class_register.values())

    @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):
            cls.verify_artifact_dict(d)
            return cls._class_register[d.pop("type")](**d)

        return d

    @classmethod
    def from_dict(cls, d):
        cls.verify_artifact_dict(d)
        return cls._recursive_load(d)

    @classmethod
    def load(cls, path, artifact_identifier=None):
        d = load_json(path)
        new_artifact = cls.from_dict(d)
        new_artifact.artifact_identifier = artifact_identifier
        return new_artifact

    def prepare(self):
        pass

    def verify(self):
        pass

    @final
    def __pre_init__(self, **kwargs):
        self._init_dict = get_raw(kwargs)

    @final
    def __post_init__(self):
        self.type = self.register_class(self.__class__)

        for field in fields(self):
            if issubtype(
                field.type, Union[Artifact, List[Artifact], Dict[str, Artifact]]
            ):
                value = getattr(self, field.name)
                value = map_values_in_place(value, maybe_recover_artifact)
                setattr(self, field.name, value)

        self.prepare()
        self.verify()

    def _to_raw_dict(self):
        return {"type": self.type, **self._init_dict}

    def save(self, path):
        data = self.to_dict()
        save_json(path, data)


def get_raw(obj):
    if isinstance(obj, Artifact):
        return obj._to_raw_dict()

    if isinstance(obj, tuple) and hasattr(obj, "_fields"):  # named tuple
        return type(obj)(*[get_raw(v) for v in obj])

    if isinstance(obj, (list, tuple)):
        return type(obj)([get_raw(v) for v in obj])

    if isinstance(obj, dict):
        return type(obj)({get_raw(k): get_raw(v) for k, v in obj.items()})

    return deepcopy(obj)


class ArtifactList(list, Artifact):
    def prepare(self):
        for artifact in self:
            artifact.prepare()


class Artifactory(Artifact):
    @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}"


@lru_cache(maxsize=None)
def fetch_artifact(name):
    if Artifact.is_artifact_file(name):
        return Artifact.load(name), None

    for artifactory in Artifactories():
        if name in artifactory:
            return artifactory[name], artifactory

    raise UnitxtArtifactNotFoundError(name, Artifactories().artifactories)


@lru_cache(maxsize=None)
def verbosed_fetch_artifact(identifer):
    artifact, artifactory = fetch_artifact(identifer)
    logger.info(f"Artifact {identifer} is fetched from {artifactory}")
    return artifact


def reset_artifacts_cache():
    fetch_artifact.cache_clear()
    verbosed_fetch_artifact.cache_clear()


def maybe_recover_artifact(artifact):
    if isinstance(artifact, str):
        return verbosed_fetch_artifact(artifact)

    return artifact


def register_all_artifacts(path):
    for loader, module_name, _is_pkg in pkgutil.walk_packages(path):
        logger.info(__name__)
        if module_name == __name__:
            continue
        logger.info(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:
                    logger.info(obj)