mypy types

[nix-pin-deps] / nix_pin_deps.py

diff --git a/nix_pin_deps.py b/nix_pin_deps.py
index 20988f58289b8a002c40f6c1a16d9c242a6357aa..2b88634f687f97044170bb9e72c799cc77fa2836 100644 (file)
--- a/nix_pin_deps.py
+++ b/nix_pin_deps.py
@@ -6,9 +6,11 @@ import sys
 import xml.sax
 import xml.sax.handler
 
+from typing import Any, Dict, Iterator, List, Set
+
 
 @contextmanager
-def log(msg):
+def log(msg: str) -> Iterator[None]:
     print(msg, file=sys.stderr, end='', flush=True)
     try:
         yield
@@ -17,12 +19,12 @@ def log(msg):
 
 
 class ParseNixStoreQueryGraphML(xml.sax.handler.ContentHandler):
-    def __init__(self):
-        self.roots = set()
-        self.non_roots = set()
-        self.deps = {}
+    def __init__(self) -> None:
+        self.roots: Set[str] = set()
+        self.non_roots: Set[str] = set()
+        self.deps: Dict[str, List[str]] = {}
 
-    def startElement(self, name, attrs):
+    def startElement(self, name: str, attrs: Any) -> None:
         if name == "edge":
             source = attrs.getValue("source")
             target = attrs.getValue("target")
@@ -34,17 +36,18 @@ class ParseNixStoreQueryGraphML(xml.sax.handler.ContentHandler):
                 self.roots.add(source)
 
 
-def getDeps(drv):
+def getDeps(drv: str) -> ParseNixStoreQueryGraphML:
     with log("Loading dependency tree..."):
         process = subprocess.Popen(
             ["nix-store", "--query", "--graphml", drv], stdout=subprocess.PIPE)
         parser = ParseNixStoreQueryGraphML()
+        assert process.stdout
         xml.sax.parse(process.stdout, parser)
         assert process.wait() == 0
         return parser
 
 
-def getDrvInfo(drv):
+def getDrvInfo(drv: str) -> Any:
     with log("Loading %s..." % drv):
         process = subprocess.Popen(["nix",
                                     "--experimental-features",
@@ -52,12 +55,13 @@ def getDrvInfo(drv):
                                     "show-derivation",
                                     "/nix/store/" + drv],
                                    stdout=subprocess.PIPE)
+        assert process.stdout
         info = json.load(process.stdout)
         assert process.wait() == 0
         return info
 
 
-def allBuilt(drv):
+def allBuilt(drv: str) -> bool:
     # TODO: How to pin outputs one at a time?
     # Currently, we only pin when all outputs are available.
     # It would be better to pin the outputs we've got.
@@ -65,15 +69,15 @@ def allBuilt(drv):
         drv).values() for o in d['outputs'].values())
 
 
-def isDrv(thing):
+def isDrv(thing: str) -> bool:
     return thing.endswith(".drv")
 
 
-def removesuffix(s, suf):
+def removesuffix(s: str, suf: str) -> str:
     return s[:-len(suf)] if s.endswith(suf) else s
 
 
-def pin(drv):
+def pin(drv: str) -> None:
     outPath = os.path.join(sys.argv[2], removesuffix(drv, ".drv"))
     if not os.path.exists(outPath):
         process = subprocess.run(["nix-store",
@@ -84,7 +88,9 @@ def pin(drv):
                                  check=True)
 
 
-def pinBuiltThings(thing, done, deps):
+def pinBuiltThings(thing: str,
+                   done: Set[str],
+                   deps: Dict[str, List[str]]) -> None:
     if thing in done:
         return
     done.add(thing)