File: variable_store.py

package info (click to toggle)
python-moto 5.1.18-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 116,520 kB
  • sloc: python: 636,725; javascript: 181; makefile: 39; sh: 3
file content (129 lines) | stat: -rw-r--r-- 4,933 bytes parent folder | download
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
from __future__ import annotations

from typing import Any, Final, Optional

from moto.stepfunctions.parser.asl.jsonata.jsonata import (
    VariableDeclarations,
    encode_jsonata_variable_declarations,
)
from moto.stepfunctions.parser.asl.utils.encoding import to_json_str

VariableIdentifier = str
VariableValue = Any


class VariableStoreError(RuntimeError):
    message: Final[str]

    def __init__(self, message: str):
        self.message = message

    def __str__(self):
        return f"{self.__class__.__name__} {self.message}"

    def __repr__(self):
        return str(self)


class NoSuchVariable(VariableStoreError):
    variable_identifier: Final[VariableIdentifier]

    def __init__(self, variable_identifier: VariableIdentifier):
        super().__init__(message=f"No such variable '{variable_identifier}' in scope")
        self.variable_identifier = variable_identifier


class IllegalOuterScopeWrite(VariableStoreError):
    variable_identifier: Final[VariableIdentifier]
    variable_value: Final[VariableValue]

    def __init__(
        self, variable_identifier: VariableIdentifier, variable_value: VariableValue
    ):
        super().__init__(
            message=f"Cannot bind value '{variable_value}' to variable '{variable_identifier}' as it belongs to an outer scope."
        )
        self.variable_identifier = variable_identifier
        self.variable_value = variable_value


class VariableStore:
    _outer_scope: Final[dict]
    _inner_scope: Final[dict]

    _declaration_tracing: Final[set[str]]

    _outer_variable_declaration_cache: Optional[VariableDeclarations]
    _variable_declarations_cache: Optional[VariableDeclarations]

    def __init__(self):
        self._outer_scope = {}
        self._inner_scope = {}
        self._declaration_tracing = set()
        self._outer_variable_declaration_cache = None
        self._variable_declarations_cache = None

    @classmethod
    def as_inner_scope_of(cls, outer_variable_store: VariableStore) -> VariableStore:
        inner_variable_store = cls()
        inner_variable_store._outer_scope.update(outer_variable_store._outer_scope)
        inner_variable_store._outer_scope.update(outer_variable_store._inner_scope)
        return inner_variable_store

    def reset_tracing(self) -> None:
        self._declaration_tracing.clear()

    # TODO: add typing when this available in service init.
    def get_assigned_variables(self) -> dict[str, str]:
        assigned_variables: dict[str, str] = {}
        for traced_declaration_identifier in self._declaration_tracing:
            traced_declaration_value = self.get(traced_declaration_identifier)
            if isinstance(traced_declaration_value, str):
                traced_declaration_value_json_str = f'"{traced_declaration_value}"'
            else:
                traced_declaration_value_json_str: str = to_json_str(
                    traced_declaration_value, separators=(",", ":")
                )
            assigned_variables[traced_declaration_identifier] = (
                traced_declaration_value_json_str
            )
        return assigned_variables

    def get(self, variable_identifier: VariableIdentifier) -> VariableValue:
        if variable_identifier in self._inner_scope:
            return self._inner_scope[variable_identifier]
        if variable_identifier in self._outer_scope:
            return self._outer_scope[variable_identifier]
        raise NoSuchVariable(variable_identifier=variable_identifier)

    def set(
        self, variable_identifier: VariableIdentifier, variable_value: VariableValue
    ) -> None:
        if variable_identifier in self._outer_scope:
            raise IllegalOuterScopeWrite(
                variable_identifier=variable_identifier, variable_value=variable_value
            )
        self._declaration_tracing.add(variable_identifier)
        self._inner_scope[variable_identifier] = variable_value
        self._variable_declarations_cache = None

    @staticmethod
    def _to_variable_declarations(bindings: dict[str, Any]) -> VariableDeclarations:
        variables = {f"${key}": value for key, value in bindings.items()}
        encoded = encode_jsonata_variable_declarations(variables)
        return encoded

    def get_variable_declarations(self) -> VariableDeclarations:
        if self._variable_declarations_cache is not None:
            return self._variable_declarations_cache
        if self._outer_variable_declaration_cache is None:
            self._outer_variable_declaration_cache = self._to_variable_declarations(
                self._outer_scope
            )
        inner_variable_declarations_cache = self._to_variable_declarations(
            self._inner_scope
        )
        self._variable_declarations_cache = "".join(
            [self._outer_variable_declaration_cache, inner_variable_declarations_cache]
        )
        return self._variable_declarations_cache