File size: 1,513 Bytes
128b1df
6754ccb
128b1df
c8ed6c2
 
128b1df
7c32c70
6754ccb
 
 
128b1df
 
6754ccb
128b1df
 
 
 
6754ccb
 
 
128b1df
6754ccb
e3813ff
 
 
7c32c70
 
 
 
 
 
 
 
 
6754ccb
 
128b1df
6754ccb
 
 
 
 
 
7c32c70
 
 
 
6754ccb
 
 
 
 
 
 
 
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
import typing
from dataclasses import field
from typing import Dict, List

from .artifact import Artifact
from .dataclass import AbstractField
from .random_utils import random


class Collection(Artifact):
    items: typing.Collection = AbstractField()

    def __getitem__(self, key):
        try:
            return self.items[key]
        except LookupError:
            raise LookupError(f"Cannot find item {repr(key)} in {repr(self)}")


class ListCollection(Collection):
    items: List[Artifact] = field(default_factory=list)

    def __len__(self):
        return len(self.items)

    def append(self, item):
        self.items.append(item)

    def extend(self, other):
        self.items.extend(other.items)

    def __add__(self, other):
        return ListCollection(self.items.__add__(other.items))


class DictCollection(Collection):
    items: Dict[str, Artifact] = field(default_factory=dict)


class ItemPicker(Artifact):
    item: object = None

    def __call__(self, collection: Collection):
        try:
            return collection[int(self.item)]
        except (SyntaxError, KeyError, ValueError) as e:  # in case picking from a dictionary
            return collection[self.item]


class RandomPicker(Artifact):
    def __call__(self, collection: Collection):
        if isinstance(collection, ListCollection):
            return random.choice(list(collection.items))
        elif isinstance(collection, DictCollection):
            return random.choice(list(collection.items.values()))