support url source types

bin/poetry2nix

@@ -7,35 +7,84 @@ import toml
 import json
 import sys
 
-from typing import Dict, Any, Tuple
+from typing import Dict, Any, Tuple, List
 
 
-def fetch_git(pkg: Dict[str, Any]) -> Tuple[str, subprocess.CompletedProcess]:
-    try:
-        reference = pkg["source"]["resolved_reference"]
-    except KeyError:
-        reference = pkg["source"]["reference"]
+class Package:
+    def __init__(self, attrs: Dict[str, Any]) -> None:
+        self.attrs = attrs
+        self.name = attrs["name"]
+        self.source = self.attrs["source"]
 
-    return (
-        pkg["name"],
-        subprocess.run(
-            [
-                "nix-prefetch-git",
-                "--fetch-submodules",
-                "--url",
-                pkg["source"]["url"],
-                "--rev",
-                reference,
-            ],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        ),
-    )
+    def fetch(self) -> Tuple["Package", subprocess.CompletedProcess]:
+        raise NotImplementedError()
+
+    def expression(self, output: str) -> str:
+        raise NotImplementedError()
 
 
-def indent(expr: str, spaces: int = 2) -> str:
-    i = " " * spaces
-    return "\n".join([(i if l != "" else "") + l for l in expr.split("\n")])
+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:
@@ -54,22 +103,31 @@ def parse_args() -> argparse.Namespace:
     return argparser.parse_args()
 
 
+def indent(expr: str, spaces: int = 2) -> str:
+    i = " " * spaces
+    return "\n".join([(i if l != "" else "") + l for l in expr.split("\n")])
+
+
 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 }:",
@@ -77,30 +135,13 @@ def main() -> None:
         ]
 
         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(["", "}", ""])
 

mk-poetry-dep.nix

@@ -46,6 +46,7 @@ pythonPackages.callPackage
       toPath = s: pwd + "/${s}";
       isSource = source != null;
       isGit = isSource && source.type == "git";
+      isUrl = isSource && source.type == "url";
       isLocal = isSource && source.type == "directory";
       localDepPath = toPath source.url;
 
@@ -95,7 +96,7 @@ pythonPackages.callPackage
         lib.optional (! lib.elem name skipSetupToolsSCM) pythonPackages.setuptools-scm
         ++ lib.optional (! lib.elem name [ "pip" "setuptools" "wheel" ]) pythonPackages.wheel
       ;
-      format = if isLocal then "pyproject" else if isGit then "pyproject" else fileInfo.format;
+      format = if isLocal || isGit || isUrl then "pyproject" else fileInfo.format;
     in
     buildPythonPackage {
       pname = moduleName name;
@@ -165,11 +166,19 @@ pythonPackages.callPackage
               rev = source.resolved_reference or source.reference;
               ref = sourceSpec.branch or sourceSpec.rev or sourceSpec.tag or "HEAD";
             }
-          ) else if isLocal then (poetryLib.cleanPythonSources { src = localDepPath; }) else
-        fetchFromPypi {
-          pname = name;
-          inherit (fileInfo) file hash kind;
-        };
+          )
+        else if isUrl then
+          builtins.fetchTarball
+            {
+              inherit (source) url;
+            }
+        else if isLocal then
+          (poetryLib.cleanPythonSources { src = localDepPath; })
+        else
+          fetchFromPypi {
+            pname = name;
+            inherit (fileInfo) file hash kind;
+          };
     }
   )
 { }