polar: minor improvements, konsume: more type annotations

my/kython/konsume.py

@@ -11,7 +11,7 @@ def zoom(w, *keys):
 
 # TODO need to support lists
 class Zoomable:
-    def __init__(self, parent, *args, **kwargs):
+    def __init__(self, parent, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs) # type: ignore
         self.parent = parent
 
@@ -21,19 +21,19 @@ class Zoomable:
     def dependants(self):
         raise NotImplementedError
 
-    def ignore(self):
+    def ignore(self) -> None:
         self.consume_all()
 
-    def consume_all(self):
+    def consume_all(self) -> None:
         for d in self.dependants:
             d.consume_all()
         self.consume()
 
-    def consume(self):
+    def consume(self) -> None:
         assert self.parent is not None
         self.parent._remove(self)
 
-    def zoom(self):
+    def zoom(self) -> 'Zoomable':
         self.consume()
         return self
 
@@ -56,6 +56,8 @@ class Wdict(Zoomable, OrderedDict):
 
     def this_consumed(self):
         return len(self) == 0
+    # TODO specify mypy type for the index special method?
+
 
 class Wlist(Zoomable, list):
     def _remove(self, xx):
@@ -83,7 +85,8 @@ class Wvalue(Zoomable):
     def __repr__(self):
         return 'WValue{' + repr(self.value) + '}'
 
-def _wrap(j, parent=None):
+from typing import Tuple
+def _wrap(j, parent=None) -> Tuple[Zoomable, List[Zoomable]]:
     res: Zoomable
     cc: List[Zoomable]
     if isinstance(j, dict):
@@ -109,13 +112,14 @@ def _wrap(j, parent=None):
         raise RuntimeError(f'Unexpected type: {type(j)} {j}')
 
 from contextlib import contextmanager
+from typing import Iterator
 
 class UnconsumedError(Exception):
     pass
 
 # TODO think about error policy later...
 @contextmanager
-def wrap(j, throw=True):
+def wrap(j, throw=True) -> Iterator[Zoomable]:
     w, children = _wrap(j)
 
     yield w
@@ -128,28 +132,33 @@ def wrap(j, throw=True):
                 # TODO log?
                 pass
 
-
+from typing import cast
 def test_unconsumed():
     import pytest # type: ignore
     with pytest.raises(UnconsumedError):
         with wrap({'a': 1234}) as w:
+            w = cast(Wdict, w)
             pass
 
     with pytest.raises(UnconsumedError):
         with wrap({'c': {'d': 2222}}) as w:
+            w = cast(Wdict, w)
             d = w['c']['d'].zoom()
 
 def test_consumed():
     with wrap({'a': 1234}) as w:
+        w = cast(Wdict, w)
         a = w['a'].zoom()
 
     with wrap({'c': {'d': 2222}}) as w:
+        w = cast(Wdict, w)
         c = w['c'].zoom()
         d = c['d'].zoom()
 
 def test_types():
     # (string, number, object, array, boolean or nul
     with wrap({'string': 'string', 'number': 3.14, 'boolean': True, 'null': None, 'list': [1, 2, 3]}) as w:
+        w = cast(Wdict, w)
         w['string'].zoom()
         w['number'].consume()
         w['boolean'].zoom()
@@ -159,5 +168,8 @@ def test_types():
 
 def test_consume_all():
     with wrap({'aaa': {'bbb': {'hi': 123}}}) as w:
+        w = cast(Wdict, w)
         aaa = w['aaa'].zoom()
         aaa['bbb'].consume_all()
+
+# TODO type check this...