Fix errors caught by type-checker.

1 files changed, 31 insertions, 26 deletions

diff --git a/gn_libs/privileges.py b/gn_libs/privileges.py
index ad813da..32c943d 100644
--- a/gn_libs/privileges.py
+++ b/gn_libs/privileges.py
@@ -1,7 +1,7 @@
 """Utilities for handling privileges."""
 import logging
 from functools import reduce
-from typing import Union, Sequence, Iterator, TypeAlias
+from typing import Union, Sequence, Iterator, TypeAlias, TypedDict
 
 logger = logging.getLogger(__name__)
 
@@ -27,16 +27,25 @@ def __add_leaves__(
         index: int,
         tree: tuple[Operator],
         leaves: dict
-) -> Union[tuple[Operator], ParseTree]:
+) -> Union[tuple[Operator], Union[ParseTree, tuple]]:
     """Add leaves to the tree."""
     if leaves.get(index):
         return tree + (leaves[index],)
-    return tree
+    return tree + (tuple())
 
 
-def __build_tree__(tree_state: dict) -> ParseTree:
+class ParsingState(TypedDict):
+    """Class to create a state object. Mostly used to silence MyPy"""
+    tokens: list[str]
+    trees: list[tuple[int, int, str, int, int]]#(name, parent, operator, start, end)
+    open_parens: int
+    current_tree: int
+    leaves: dict[int, tuple[str, ...]]#[parent-tree, [index, index, ...]]
+
+
+def __build_tree__(tree_state: ParsingState) -> ParseTree:
     """Given computed state, build the actual tree."""
-    _built= []
+    _built = []
     for idx, tree in enumerate(tree_state["trees"]):
         _built.append(__add_leaves__(idx, (tree[2],), tree_state["leaves"]))
 
@@ -70,49 +79,45 @@ def __build_tree__(tree_state: dict) -> ParseTree:
 
 def __parse_tree__(tokens: Iterator[str]) -> ParseTree:
     """Parse the tokens into a tree."""
-    _state = {
-        "tokens": [],
-        "trees": [],#(name, parent, operator, start, end)
-        "open-parens": 0,
-        "current-tree": 0,
-        "leaves": {}#[parent-tree, [index, index, ...]]
-    }
+    _state = ParsingState(
+        tokens=[], trees=[], open_parens=0, current_tree=0, leaves={})
     for _idx, _token in enumerate(tokens):
         _state["tokens"].append(_token)
 
         if _idx==0:
             if _token[1:].upper() not in _OPERATORS_:
                 raise SpecificationValueError(f"Invalid operator: {_token[1:]}")
-            _state["open-parens"] += 1
-            _state["trees"].append((0, 0, _token[1:].upper(), _idx))
-            _state["current-tree"] = 0
+            _state["open_parens"] += 1
+            _state["trees"].append((0, 0, _token[1:].upper(), _idx, -1))
+            _state["current_tree"] = 0
             continue# this is bad!
 
         if _token == ")":# end a tree
             logger.debug("ENDING A TREE: %s", _state)
-            _state["open-parens"] -= 1
-            _state["trees"][_state["current-tree"]] = (
-                _state["trees"][_state["current-tree"]] + (_idx,))
+            _state["open_parens"] -= 1
+            _state["trees"][_state["current_tree"]] = (
+                _state["trees"][_state["current_tree"]][0:-1] + (_idx,))
             # We go back to the parent below.
-            _state["current-tree"] = _state["trees"][_state["current-tree"]][1]
+            _state["current_tree"] = _state["trees"][_state["current_tree"]][1]
             continue# still really bad!
 
         if _token[1:].upper() in _OPERATORS_:# new child tree
-            _state["open-parens"] += 1
+            _state["open_parens"] += 1
             _state["trees"].append((len(_state["trees"]),
-                                    _state["current-tree"],
+                                    _state["current_tree"],
                                     _token[1:].upper(),
-                                    _idx))
-            _state["current-tree"] = len(_state["trees"]) - 1
+                                    _idx,
+                                    -1))
+            _state["current_tree"] = len(_state["trees"]) - 1
             continue# more evil still
 
         logger.debug("state: %s", _state)
         # leaves
-        _state["leaves"][_state["current-tree"]] = _state["leaves"].get(
-            _state["current-tree"], tuple()) + (_token,)
+        _state["leaves"][_state["current_tree"]] = _state["leaves"].get(
+            _state["current_tree"], tuple()) + (_token,)
 
     # Build parse-tree from state
-    if _state["open-parens"] != 0:
+    if _state["open_parens"] != 0:
         raise SpecificationValueError("Unbalanced parentheses.")
     return __build_tree__(_state)