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

diff --git a/nix_pin_deps.py b/nix_pin_deps.py
index 2b88634..6f12df8 100644
--- a/nix_pin_deps.py
+++ b/nix_pin_deps.py
@@ -20,6 +20,7 @@ def log(msg: str) -> Iterator[None]:
 
 class ParseNixStoreQueryGraphML(xml.sax.handler.ContentHandler):
     def __init__(self) -> None:
+        super().__init__()
         self.roots: Set[str] = set()
         self.non_roots: Set[str] = set()
         self.deps: Dict[str, List[str]] = {}
@@ -38,26 +39,24 @@ class ParseNixStoreQueryGraphML(xml.sax.handler.ContentHandler):
 
 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
+        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: str) -> Any:
-    with log("Loading %s..." % drv):
-        process = subprocess.Popen(["nix",
-                                    "--experimental-features",
-                                    "nix-command",
-                                    "show-derivation",
-                                    "/nix/store/" + drv],
-                                   stdout=subprocess.PIPE)
-        assert process.stdout
-        info = json.load(process.stdout)
-        assert process.wait() == 0
+    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
 
 
@@ -80,12 +79,8 @@ def removesuffix(s: str, suf: str) -> str:
 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: str,
@@ -101,6 +96,11 @@ def pinBuiltThings(thing: str,
         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()