1 files changed, 112 insertions, 58 deletions
diff --git a/infra/libkookie/nixpkgs/pkgs/development/tools/poetry2nix/poetry2nix/bin/poetry2nix b/infra/libkookie/nixpkgs/pkgs/development/tools/poetry2nix/poetry2nix/bin/poetry2nix
index 355cebfd50c4..559025f3ec45 100755
--- a/infra/libkookie/nixpkgs/pkgs/development/tools/poetry2nix/poetry2nix/bin/poetry2nix
+++ b/infra/libkookie/nixpkgs/pkgs/development/tools/poetry2nix/poetry2nix/bin/poetry2nix
@@ -7,60 +7,127 @@ import toml
 import json
 import sys
 
-
-argparser = argparse.ArgumentParser(description="Poetry2nix CLI")
-
-subparsers = argparser.add_subparsers(dest="subcommand")
-subparsers.required = True
-
-parser_lock = subparsers.add_parser("lock", help="Generate overrides for git hashes",)
-parser_lock.add_argument(
-    "--lock", default="poetry.lock", help="Path to input poetry.lock",
-)
-parser_lock.add_argument(
-    "--out", default="poetry-git-overlay.nix", help="Output file",
-)
-
-
-def fetch_git(pkg):
-    return (
-        pkg["name"],
-        subprocess.run(
-            [
-                "nix-prefetch-git",
-                "--fetch-submodules",
-                "--url",
-                pkg["source"]["url"],
-                "--rev",
-                pkg["source"]["reference"],
-            ],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        ),
+from typing import Dict, Any, Tuple, List
+
+
+class Package:
+    def __init__(self, attrs: Dict[str, Any]) -> None:
+        self.attrs = attrs
+        self.name = attrs["name"]
+        self.source = self.attrs["source"]
+
+    def fetch(self) -> Tuple["Package", subprocess.CompletedProcess]:
+        raise NotImplementedError()
+
+    def expression(self, output: str) -> str:
+        raise NotImplementedError()
+
+
+class UrlPackage(Package):
+    def fetch(self) -> Tuple[Package, subprocess.CompletedProcess]:
+        return (
+            self,
+            subprocess.run(
+                [
+                    "nix-prefetch-url",
+                    "--unpack",
+                    self.source["url"],
+                ],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True
+            ),
+        )
+
+    def expression(self, output: str) -> str:
+        sha256 = output.rstrip()
+        return textwrap.dedent("""
+            %s = super.%s.overridePythonAttrs (
+              _: {
+                src = pkgs.fetchzip {
+                  url = "%s";
+                  sha256 = "%s";
+                };
+              }
+            );""" % (self.name, self.name, self.source["url"], sha256))
+
+
+class GitPackage(Package):
+    def fetch(self) -> Tuple[Package, subprocess.CompletedProcess]:
+        reference = self.source.get("resolved_reference", self.source["reference"])
+
+        return (
+            self,
+            subprocess.run(
+                [
+                    "nix-prefetch-git",
+                    "--fetch-submodules",
+                    "--url",
+                    self.source["url"],
+                    "--rev",
+                    reference,
+                ],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True
+            ),
+        )
+
+    def expression(self, output: str) -> str:
+        meta = json.loads(output)
+        return textwrap.dedent("""
+            %s = super.%s.overridePythonAttrs (
+              _: {
+                src = pkgs.fetchgit {
+                  url = "%s";
+                  rev = "%s";
+                  sha256 = "%s";
+                };
+              }
+            );""" % (self.name, self.name, meta["url"], meta["rev"], meta["sha256"]))
+
+
+def parse_args() -> argparse.Namespace:
+    argparser = argparse.ArgumentParser(description="Poetry2nix CLI")
+
+    subparsers = argparser.add_subparsers(dest="subcommand")
+    subparsers.required = True
+
+    parser_lock = subparsers.add_parser("lock", help="Generate overrides for git hashes",)
+    parser_lock.add_argument(
+        "--lock", default="poetry.lock", help="Path to input poetry.lock",
+    )
+    parser_lock.add_argument(
+        "--out", default="poetry-git-overlay.nix", help="Output file",
     )
+    return argparser.parse_args()
 
 
-def indent(expr, spaces=2):
+def indent(expr: str, spaces: int = 2) -> str:
     i = " " * spaces
     return "\n".join([(i if l != "" else "") + l for l in expr.split("\n")])
 
 
-if __name__ == "__main__":
-    args = argparser.parse_args()
+def main() -> None:
+    args = parse_args()
 
     with open(args.lock) as lockf:
         lock = toml.load(lockf)
 
-    pkgs = []
+    pkgs: List[Package] = []
     for pkg in lock["package"]:
         if "source" in pkg:
-            pkgs.append(pkg)
+            source_type = pkg["source"]["type"]
+            if source_type == "git":
+                pkgs.append(GitPackage(pkg))
+            elif source_type == "url":
+                pkgs.append(UrlPackage(pkg))
 
     with ThreadPoolExecutor() as e:
         futures = []
 
         for pkg in pkgs:
-            futures.append(e.submit(fetch_git, pkg))
+            futures.append(e.submit(pkg.fetch))
 
         lines = [
             "{ pkgs }:",
@@ -68,30 +135,13 @@ if __name__ == "__main__":
         ]
 
         for f in futures:
-            drv_name, p = f.result()
+            package, p = f.result()
             if p.returncode != 0:
-                sys.stderr.buffer.write(p.stderr)
-                sys.stderr.buffer.flush()
+                sys.stderr.write(p.stderr)
+                sys.stderr.flush()
                 exit(p.returncode)
-
-            meta = json.loads(p.stdout.decode())
-            lines.append(
-                indent(
-                    textwrap.dedent(
-                        """
-              %s = super.%s.overridePythonAttrs (
-                _: {
-                  src = pkgs.fetchgit {
-                    url = "%s";
-                    rev = "%s";
-                    sha256 = "%s";
-                  };
-                }
-              );"""
-                        % (drv_name, drv_name, meta["url"], meta["rev"], meta["sha256"])
-                    )
-                )
-            )
+            expr = package.expression(p.stdout)
+            lines.append(indent(expr))
 
         lines.extend(["", "}", ""])
 
@@ -101,3 +151,7 @@ if __name__ == "__main__":
         fout.write(expr)
 
     print(f"Wrote {args.out}")
+
+
+if __name__ == "__main__":
+    main()