h2o-llmstudio

1
from typing import Iterable, List, Optional
2

3

4
class Order:
5
    """
6
    A list-like structure to specify the order of items in a dictionary.
7
    The main characteristics are:
8
        - Append and insert only. Cannot remove elements. This is not strictly required
9
            by the use-case but probably good practice.
10
        - Elements must be unique. Inserting an element which is already in the list
11
            will throw an error.
12

13
    Primarily useful for specifying the order in which UI elements
14
    should be shown in Wave.
15
    """
16

17
    def __init__(self, keys: Optional[List[str]] = None):
18
        if keys is not None:
19
            self._list = list(keys)
20
        else:
21
            self._list = list()
22

23
    def _unique_guard(self, *keys: str):
24
        for key in keys:
25
            if key in self._list:
26
                raise ValueError(f"`{key}` is already in the list!")
27

28
    def append(self, key: str):
29
        """
30
        Append a key at the end of the list:
31

32
        Args:
33
            key: String to append.
34

35
        Raises:
36
            - `ValueError` if the key is already in the list.
37
        """
38

39
        self._unique_guard(key)
40

41
        self._list.append(key)
42

43
    def extend(self, keys: Iterable[str]):
44
        """
45
        Extend the list by multiple keys.
46

47
        Args:
48
            keys: Iterable of keys.
49

50
        Raises:
51
            - `ValueError` if one or more key is already in the list.
52
        """
53

54
        self._unique_guard(*keys)
55

56
        self._list.extend(keys)
57

58
    def insert(
59
        self, *keys: str, before: Optional[str] = None, after: Optional[str] = None
60
    ):
61
        """
62
        Insert one or more keys. Either `before` or `after`, but not both, must be set
63
        to determine position.
64

65
        Args:
66
            keys: One more keys to insert.
67
            after: A key immediately after which the `keys` will be inserted.
68
            before: A key immediately before which the `keys` are inserted.
69

70
        Raises:
71
            - `ValueError` if one or more key is already in the list.
72
            - `ValueError` if `before` / `after` does not exist in the list.
73
            - `ValueError` if an invalid combination of arguments is set.
74
        """
75

76
        self._unique_guard(*keys)
77

78
        if before is not None:
79
            for key in keys[::-1]:
80
                self._list.insert(self._list.index(before), key)
81

82
            if after is not None:
83
                raise ValueError("`after` must be None if `before` is set.")
84

85
        if after is not None:
86
            for key in keys[::-1]:
87
                self._list.insert(self._list.index(after) + 1, key)
88

89
            if before is not None:
90
                raise ValueError("`before` must be None if `after` is set.")
91

92
        if before is None and after is None:
93
            raise ValueError("Either `before` or `after` must be set.")
94

95
    def __getitem__(self, idx):
96
        return self._list[idx]
97

98
    def __len__(self):
99
        return len(self._list)
100

101
    def __iter__(self):
102
        return iter(self._list)
103

104

105
def test_order():
106
    order = Order(["dataset", "training", "validation", "logging"])
107

108
    order.insert("architecture", before="training")
109
    order.insert("environment", after="validation")
110

111
    assert [item for item in order] == [
112
        "dataset",
113
        "architecture",
114
        "training",
115
        "validation",
116
        "environment",
117
        "logging",
118
    ]
119