X-Git-Url: http://git.scottworley.com/nix-pin-deps/blobdiff_plain/1e40791b349d84c9f9398809dc4d94e832ac4725..107c5850403d88393a819538426fe9594f5746ec:/nix_pin_deps.py?ds=inline

diff --git a/nix_pin_deps.py b/nix_pin_deps.py
index 20988f5..6f12df8 100644
--- 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,13 @@ 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:
+        super().__init__()
+        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,30 +37,30 @@ 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()
-        xml.sax.parse(process.stdout, parser)
-        assert process.wait() == 0
+        with subprocess.Popen(
+                ["nix-store", "--query", "--graphml", drv], stdout=subprocess.PIPE) as process:
+            parser = ParseNixStoreQueryGraphML()
+            assert process.stdout
+            xml.sax.parse(process.stdout, parser)
+            assert process.wait() == 0
         return parser
 
 
-def getDrvInfo(drv):
-    with log("Loading %s..." % drv):
-        process = subprocess.Popen(["nix",
-                                    "--experimental-features",
-                                    "nix-command",
-                                    "show-derivation",
-                                    "/nix/store/" + drv],
-                                   stdout=subprocess.PIPE)
-        info = json.load(process.stdout)
-        assert process.wait() == 0
+def getDrvInfo(drv: str) -> Any:
+    with log(f"Loading {drv}..."):
+        with subprocess.Popen(
+                ["nix", "--experimental-features", "nix-command",
+                    "show-derivation", "/nix/store/" + drv],
+                stdout=subprocess.PIPE) as process:
+            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,26 +68,24 @@ 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",
-                                  "--realise",
-                                  "--add-root",
-                                  outPath,
-                                  "/nix/store/" + drv],
-                                 check=True)
+        subprocess.run(["nix-store", "--realise", "--add-root",
+                       outPath, "/nix/store/" + 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)
@@ -95,6 +96,11 @@ def pinBuiltThings(thing, done, deps):
         pinBuiltThings(dep, done, deps)
 
 
-dep_graph = getDeps(sys.argv[1])
-for root in dep_graph.roots:
-    pinBuiltThings(root, set(), dep_graph.deps)
+def main() -> None:
+    dep_graph = getDeps(sys.argv[1])
+    for root in dep_graph.roots:
+        pinBuiltThings(root, set(), dep_graph.deps)
+
+
+if __name__ == '__main__':
+    main()