Adding type annotations

.gitignore

@@ -2,3 +2,4 @@
 *.sw*
 micro-benchmark/*_test.py
 micro-benchmark-key-errs/*_test.py
+.direnv/

pycg/__main__.py

@@ -1,13 +1,14 @@
+import argparse
+import json
 import os
 import sys
-import json
-import argparse
 
-from pycg.pycg import CallGraphGenerator
 from pycg import formats
+from pycg.pycg import CallGraphGenerator
 from pycg.utils.constants import CALL_GRAPH_OP, KEY_ERR_OP
 
-def main():
+
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("entry_point",
         nargs="*",

pycg/formats/__init__.py

@@ -18,7 +18,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+from .as_graph import AsGraph
 from .fasten import Fasten
-from .simple import Simple
 from .fuzz import Fuzz
-from .as_graph import AsGraph
+from .simple import Simple

pycg/formats/as_graph.py

@@ -20,6 +20,7 @@
 #
 from .base import BaseFormatter
 
+
 class AsGraph(BaseFormatter):
     def __init__(self, cg_generator):
         self.cg_generator = cg_generator

pycg/formats/fasten.py

@@ -22,12 +22,14 @@
 
 from pkg_resources import Requirement
 
+from pycg import utils
+
 from .base import BaseFormatter
 
-from pycg import utils
+from pycg.pycg import CallGraphGenerator
 
 class Fasten(BaseFormatter):
-    def __init__(self, cg_generator, package, product, forge, version, timestamp):
+    def __init__(self, cg_generator: CallGraphGenerator, package, product, forge, version, timestamp) -> None:
         self.cg_generator = cg_generator
         self.internal_mods = self.cg_generator.output_internal_mods() or {}
         self.external_mods = self.cg_generator.output_external_mods() or {}
@@ -42,12 +44,12 @@ def __init__(self, cg_generator, package, product, forge, version, timestamp):
         self.version = version
         self.timestamp = timestamp
 
-    def get_unique_and_increment(self):
+    def get_unique_and_increment(self) -> int:
         unique = self.unique
         self.unique += 1
         return unique
 
-    def to_uri(self, modname, name=""):
+    def to_uri(self, modname: str, name: str = "") -> str:
         cleared = name
         if name:
             if name == modname:
@@ -64,7 +66,7 @@ def to_uri(self, modname, name=""):
 
         return "/{}/{}{}".format(modname.replace("-", "_"), cleared, suffix)
 
-    def to_external_uri(self, modname, name=""):
+    def to_external_uri(self, modname: str, name: str = "") -> str:
         if modname == utils.constants.BUILTIN_NAME:
             name = name[len(modname)+1:]
             modname = ".builtin"

pycg/formats/fuzz.py

@@ -20,6 +20,7 @@
 #
 from .base import BaseFormatter
 
+
 class Fuzz(BaseFormatter):
     def __init__(self, cg_generator):
         self.cg_generator = cg_generator

pycg/formats/simple.py

@@ -20,6 +20,7 @@
 #
 from .base import BaseFormatter
 
+
 class Simple(BaseFormatter):
     def __init__(self, cg_generator):
         self.cg_generator = cg_generator

pycg/machinery/callgraph.py

@@ -19,21 +19,24 @@
 # under the License.
 #
 import logging
+from typing import Dict, Set, List, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
 
-class CallGraph(object):
-    def __init__(self):
-        self.cg = {}
+class CallGraph:
+    cg: Dict[str, set]
+
+    def __init__(self) -> None:
+        self.cg: Dict[str, Set[str]] = {}
         self.cg_extended = {}
-        self.modnames = {}
-        self.ep = None
-        self.entrypoints = []
+        self.modnames: Dict[str, str] = {}
+        self.ep: Optional[str] = None
+        self.entrypoints: List[Tuple[str, str]] = []
 
-        self.function_line_numbers = dict()
+        self.function_line_numbers: Dict[str, Set[int]] = {}
 
-    def add_node(self, name, modname=""):
+    def add_node(self, name: str, modname: str = ""):
         if not isinstance(name, str):
             raise CallGraphError("Only string node names allowed")
         if not name:
@@ -65,7 +68,9 @@ def add_node(self, name, modname=""):
             #else:
                 #logger.info("AN7")
 
-    def add_edge(self, src, dest, lineno=-1, mod="", ext_mod=""):
+    def add_edge(
+        self, src: str, dest: str, lineno: int = -1, mod: str = "", ext_mod: str = ""
+    ):
         self.add_node(src, mod)
         self.add_node(dest)
         self.cg[src].add(dest)
@@ -81,23 +86,23 @@ def add_edge(self, src, dest, lineno=-1, mod="", ext_mod=""):
         )
         #logger.debug(self.cg_extended[src])
 
-    def get(self):
+    def get(self) -> Dict[str, Set[str]]:
         return self.cg
 
-    def get_extended(self):
+    def get_extended(self) -> Dict:
         return self.cg_extended
 
-    def get_edges(self):
+    def get_edges(self) -> List[List[str]]:
         output = []
         for src in self.cg:
             for dst in self.cg[src]:
                 output.append([src, dst])
         return output
 
-    def get_modules(self):
+    def get_modules(self) -> Dict[str, str]:
         return self.modnames
 
-    def add_entrypoint(self, ep, modname=""):
+    def add_entrypoint(self, ep: str, modname: str = "") -> None:
         self.ep = ep
         self.ep_mod = modname
         self.entrypoints.append((ep, modname))

pycg/machinery/classes.py

@@ -18,16 +18,20 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+
+from typing import Dict, List
+
+
 class ClassManager:
-    def __init__(self):
-        self.names = {}
-        self.inheritance = {}
+    def __init__(self) -> None:
+        self.names: Dict[str, ClassNode] = {}
+        self.inheritance: Dict[str, set] = {}
 
-    def get(self, name):
+    def get(self, name: str):
         if name in self.names:
             return self.names[name]
 
-    def create(self, name, module):
+    def create(self, name: str, module: str):
         if not name in self.names:
             cls = ClassNode(name, module)
             self.names[name] = cls
@@ -36,19 +40,20 @@ def create(self, name, module):
 
         return self.names[name]
 
-    def add_inheritance(self, name, parent):
+    def add_inheritance(self, name: str, parent):
         if name not in self.inheritance:
             return
         self.inheritance[name].add(parent)
 
-    def get_classes(self):
+    def get_classes(self) -> Dict[str, "ClassNode"]:
         return self.names
 
+
 class ClassNode:
-    def __init__(self, ns, module):
+    def __init__(self, ns: str, module: str) -> None:
         self.ns = ns
         self.module = module
-        self.mro = [ns]
+        self.mro: List[str] = [ns]
 
     def add_parent(self, parent):
         if isinstance(parent, str):
@@ -59,23 +64,24 @@ def add_parent(self, parent):
                     if self.mro == parent:
                         print("This should never happen and will cause an eternal loop")
                         import sys
+
                         sys.exit(123)
 
                     self.mro.append(item)
         self.fix_mro()
 
-    def fix_mro(self):
+    def fix_mro(self) -> None:
         new_mro = []
         for idx, item in enumerate(self.mro):
-            if self.mro[idx+1:].count(item) > 0:
+            if self.mro[idx + 1 :].count(item) > 0:
                 continue
             new_mro.append(item)
         self.mro = new_mro
 
-    def get_mro(self):
+    def get_mro(self) -> List[str]:
         return self.mro
 
-    def get_module(self):
+    def get_module(self) -> str:
         return self.module
 
     def compute_mro(self):