Clean up type annotations

MarkKoz · 64json · commit 62193206d4e1 · 2019-06-11T15:50:09.000+09:00
* Add SerializableSequence alias to support tuples in addition to lists
* Replace list comprehension with a tuple
* Add missing type annotations
* Remove unused imports
diff --git a/algorithm_visualizer/commander.py b/algorithm_visualizer/commander.py
@@ -1,6 +1,6 @@
 import json
 import string
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from algorithm_visualizer import Randomize
 from algorithm_visualizer.types import Serializable, Undefined
@@ -12,9 +12,9 @@
 class Commander:
     _keyRandomizer = Randomize.String(8, string.ascii_lowercase + string.digits)
     _objectCount = 0
-    commands: List[Dict[str, Any]] = []
+    commands: List[Dict[str, Serializable]] = []
 
-    def __init__(self, *args):
+    def __init__(self, *args: Serializable):
         self._objectCount += 1
         self.key = self._keyRandomizer.create()
         self.command(self.__class__.__name__, *args)
diff --git a/algorithm_visualizer/layouts.py b/algorithm_visualizer/layouts.py
@@ -1,12 +1,12 @@
-from typing import List
+from typing import Iterable
 
 from algorithm_visualizer import Commander
 from .types import UNDEFINED
 
 
 class Layout(Commander):
-    def __init__(self, children: List[Commander]):
-        super().__init__([c.key for c in children])
+    def __init__(self, children: Iterable[Commander]):
+        super().__init__(tuple(c.key for c in children))
 
     @classmethod
     def setRoot(cls, child: Commander):
diff --git a/algorithm_visualizer/tracers/array1d.py b/algorithm_visualizer/tracers/array1d.py
@@ -1,12 +1,10 @@
-from typing import List
-
 from . import chart
 from .array2d import Array2DTracer
-from algorithm_visualizer.types import Serializable, UNDEFINED
+from algorithm_visualizer.types import Serializable, SerializableSequence, UNDEFINED
 
 
 class Array1DTracer(Array2DTracer):
-    def set(self, array1d: List[Serializable] = UNDEFINED):
+    def set(self, array1d: SerializableSequence[Serializable] = UNDEFINED):
         self.command("set", array1d)
 
     def patch(self, x: int, v: Serializable = UNDEFINED):
diff --git a/algorithm_visualizer/tracers/array2d.py b/algorithm_visualizer/tracers/array2d.py
@@ -1,11 +1,9 @@
-from typing import List
-
 from .tracer import Tracer
-from algorithm_visualizer.types import Serializable, UNDEFINED
+from algorithm_visualizer.types import Serializable, SerializableSequence, UNDEFINED
 
 
 class Array2DTracer(Tracer):
-    def set(self, array2d: List[List[Serializable]] = UNDEFINED):
+    def set(self, array2d: SerializableSequence[SerializableSequence[Serializable]] = UNDEFINED):
         self.command("set", array2d)
 
     def patch(self, x: int, y: int, v: Serializable = UNDEFINED):
diff --git a/algorithm_visualizer/tracers/graph.py b/algorithm_visualizer/tracers/graph.py
@@ -1,36 +1,34 @@
-from typing import List, Optional
-
 from .log import LogTracer, Tracer
-from algorithm_visualizer.types import Number, Serializable, UNDEFINED
+from algorithm_visualizer.types import Number, Serializable, SerializableSequence, UNDEFINED
 
 
 class GraphTracer(Tracer):
-    def set(self, array2d: List[List[Serializable]] = UNDEFINED):
+    def set(self, array2d: SerializableSequence[SerializableSequence[Serializable]] = UNDEFINED):
         self.command("set", array2d)
 
     def directed(self, isDirected: bool = UNDEFINED):
         self.command("directed", isDirected)
 
-    def weighted(self, isWeighted: bool = UNDEFINED):
+    def weighted(self, isWeighted: bool = UNDEFINED) -> "GraphTracer":
         self.command("weighted", isWeighted)
         return self
 
-    def layoutCircle(self):
+    def layoutCircle(self) -> "GraphTracer":
         self.command("layoutCircle")
         return self
 
-    def layoutTree(self, root: Serializable = UNDEFINED, sorted: bool = UNDEFINED):
+    def layoutTree(self, root: Serializable = UNDEFINED, sorted: bool = UNDEFINED) -> "GraphTracer":
         self.command("layoutTree", root, sorted)
         return self
 
-    def layoutRandom(self):
+    def layoutRandom(self) -> "GraphTracer":
         self.command("layoutRandom")
         return self
 
     def addNode(
         self,
         id: Serializable,
-        weight: Optional[Number] = UNDEFINED,
+        weight: Number = UNDEFINED,
         x: int = 0,
         y: int = 0,
         visitedCount: int = 0,
@@ -41,7 +39,7 @@ def addNode(
     def updateNode(
         self,
         id: Serializable,
-        weight: Optional[Number] = UNDEFINED,
+        weight: Number = UNDEFINED,
         x: int = UNDEFINED,
         y: int = UNDEFINED,
         visitedCount: int = UNDEFINED,
@@ -56,7 +54,7 @@ def addEdge(
         self,
         source: Serializable,
         target: Serializable,
-        weight: Optional[Number] = UNDEFINED,
+        weight: Number = UNDEFINED,
         visitedCount: int = UNDEFINED,
         selectedCount: int = UNDEFINED
     ):
@@ -66,7 +64,7 @@ def updateEdge(
         self,
         source: Serializable,
         target: Serializable,
-        weight: Optional[Number] = UNDEFINED,
+        weight: Number = UNDEFINED,
         visitedCount: int = UNDEFINED,
         selectedCount: int = UNDEFINED
     ):
@@ -79,15 +77,15 @@ def visit(
         self,
         target: Serializable,
         source: Serializable = UNDEFINED,
-        weight: Optional[Number] = UNDEFINED
+        weight: Number = UNDEFINED
     ):
         self.command("visit", target, source, weight)
 
     def leave(
         self,
         target: Serializable,
         source: Serializable = UNDEFINED,
-        weight: Optional[Number] = UNDEFINED
+        weight: Number = UNDEFINED
     ):
         self.command("leave", target, source, weight)
 
diff --git a/algorithm_visualizer/types.py b/algorithm_visualizer/types.py
@@ -1,7 +1,8 @@
 from typing import Union
 
 # Types which are serializable by the default JSONEncoder
-Serializable = Union[dict, list, str, int, float, bool, None]
+Serializable = Union[dict, list, tuple, str, int, float, bool, None]
+SerializableSequence = Union[list, tuple]
 Number = Union[int, float]
 
 
