cleanup, add/correct type hints

TUCAN-nest · Nov 15, 2023 · 5128b07 · 5128b07
1 parent 0890e4e
commit 5128b07
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/tests/parser/test_parser.py b/tests/parser/test_parser.py
@@ -45,14 +45,14 @@ def _extract_atoms_from_sum_formula(s):
         (
             "Cu",
             [
-                {ATOMIC_NUMBER: 29, ELEMENT_SYMBOL: "Cu", PARTITION: 0},
+                {ELEMENT_SYMBOL: "Cu", ATOMIC_NUMBER: 29, PARTITION: 0},
             ],
         ),
         (
             "CU",
             [
-                {ATOMIC_NUMBER: 6, ELEMENT_SYMBOL: "C", PARTITION: 0},
-                {ATOMIC_NUMBER: 92, ELEMENT_SYMBOL: "U", PARTITION: 0},
+                {ELEMENT_SYMBOL: "C", ATOMIC_NUMBER: 6, PARTITION: 0},
+                {ELEMENT_SYMBOL: "U", ATOMIC_NUMBER: 92, PARTITION: 0},
             ],
         ),
     ],

diff --git a/tucan/parser/parser.py b/tucan/parser/parser.py
@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import networkx as nx
-from typing import Any
+from typing import Any, Final
 from antlr4 import InputStream, CommonTokenStream
 from antlr4.error.ErrorListener import ErrorListener
 from antlr4.tree.Tree import ParseTreeWalker
@@ -34,7 +36,7 @@ def graph_from_tucan(tucan: str) -> nx.Graph:
     return listener.to_graph()
 
 
-def _prepare_parser(to_parse: str) -> nx.Graph:
+def _prepare_parser(to_parse: str) -> tucanParser:
     stream = InputStream(to_parse)
     lexer = tucanLexer(stream)
     token_stream = CommonTokenStream(lexer)
@@ -49,14 +51,14 @@ def _prepare_parser(to_parse: str) -> nx.Graph:
     return parser
 
 
-def _walk_tree(tree):
+def _walk_tree(tree) -> TucanListenerImpl:
     walker = ParseTreeWalker()
     listener = TucanListenerImpl()
     walker.walk(listener, tree)
     return listener
 
 
-_DESERIALIZER_NODE_ATTRIBUTE_MAPPING = {
+_DESERIALIZER_NODE_ATTRIBUTE_MAPPING: Final[dict[str, str]] = {
     v: k for k, v in _SERIALIZER_NODE_ATTRIBUTE_MAPPING.items()
 }
 

diff --git a/tucan/serialization.py b/tucan/serialization.py
@@ -10,7 +10,7 @@
 )
 from tucan.graph_utils import sort_molecule_by_attribute
 from operator import gt, lt, eq
-from typing import Callable
+from typing import Callable, Final
 import networkx as nx
 
 
@@ -34,7 +34,7 @@ def _write_edge_list(m: nx.Graph) -> str:
     return edge_list_string
 
 
-_SERIALIZER_NODE_ATTRIBUTE_MAPPING = {
+_SERIALIZER_NODE_ATTRIBUTE_MAPPING: Final[dict[str, str]] = {
     MASS: "mass",
     RAD: "rad",
 }