poetry2nix/default.nix

{ pkgs ? import <nixpkgs> {}
, lib ? pkgs.lib
, poetry ? null
, poetryLib ? import ./lib.nix { inherit lib pkgs; }
}:

let
  inherit (poetryLib) isCompatible readTOML;

  defaultPoetryOverrides = [ (import ./overrides.nix { inherit pkgs; }) ];

  mkEvalPep508 = import ./pep508.nix {
    inherit lib;
    stdenv = pkgs.stdenv;
  };

  getAttrDefault = attribute: set: default: (
    if builtins.hasAttr attribute set
    then builtins.getAttr attribute set
    else default
  );

  #
  # Returns an attrset { python, poetryPackages } for the given lockfile
  #
  mkPoetryPython =
    { poetrylock
    , poetryPkg
    , overrides ? defaultPoetryOverrides
    , meta ? {}
    , python ? pkgs.python3
    , pwd ? null
    }@attrs: let
      lockData = readTOML poetrylock;
      lockFiles = lib.getAttrFromPath [ "metadata" "files" ] lockData;

      specialAttrs = [ "poetrylock" "overrides" ];
      passedAttrs = builtins.removeAttrs attrs specialAttrs;

      evalPep508 = mkEvalPep508 python;

      # Filter packages by their PEP508 markers
      partitions = let
        supportsPythonVersion = pkgMeta: if pkgMeta ? marker then (evalPep508 pkgMeta.marker) else true;
      in
        lib.partition supportsPythonVersion lockData.package;

      compatible = partitions.right;
      incompatible = partitions.wrong;

      # Create an overriden version of pythonPackages
      #
      # We need to avoid mixing multiple versions of pythonPackages in the same
      # closure as python can only ever have one version of a dependency
      baseOverlay = self: super:
        let
          getDep = depName: if builtins.hasAttr depName self then self."${depName}" else throw "foo";

          lockPkgs = builtins.listToAttrs (
            builtins.map (
              pkgMeta: rec {
                name = pkgMeta.name;
                value = self.mkPoetryDep (
                  pkgMeta // {
                    inherit pwd;
                    source = getAttrDefault "source" pkgMeta null;
                    files = lockFiles.${name};
                    pythonPackages = self;
                  }
                );
              }
            ) compatible
          );
        in
          lockPkgs;

      overlays = [
        (
          self: super: {
            mkPoetryDep = self.callPackage ./mk-poetry-dep.nix {
              inherit pkgs lib python poetryLib;
            };
            poetry = poetryPkg;
          }
        )
        # Null out any filtered packages, we don't want python.pkgs from nixpkgs
        (self: super: builtins.listToAttrs (builtins.map (x: { name = x.name; value = null; }) incompatible))
        # Create poetry2nix layer
        baseOverlay
      ] ++ # User provided overrides
      overrides;

      packageOverrides = self: super: (
        builtins.foldl'
          (
            acc: v: let
              newSuper = acc // v self acc;
            in
              newSuper
          )
          super
          overlays
      );

      py = python.override { inherit packageOverrides; self = py; };
    in
      {
        python = py;
        poetryPackages = map (pkg: py.pkgs.${pkg.name}) compatible;
      };

  #
  # Creates a python environment with the python packages from the specified lockfile
  #
  mkPoetryEnv =
    { poetrylock
    , overrides ? defaultPoetryOverrides
    , meta ? {}
    , pwd ? null
    , python ? pkgs.python3
    }:
      let
        poetryPkg = poetry.override { inherit python; };
        py = mkPoetryPython (
          {
            inherit poetryPkg poetrylock overrides meta python pwd;
          }
        );
      in
        py.python.withPackages (_: py.poetryPackages);


  #
  # Creates a python application
  #
  mkPoetryApplication =
    { src
    , pyproject
    , poetrylock
    , overrides ? defaultPoetryOverrides
    , meta ? {}
    , python ? pkgs.python3
    , pwd ? null
    , ...
    }@attrs: let
      poetryPkg = poetry.override { inherit python; };

      py = (
        mkPoetryPython {
          inherit poetryPkg poetrylock overrides meta python pwd;
        }
      ).python;

      pyProject = readTOML pyproject;

      specialAttrs = [ "pyproject" "poetrylock" "overrides" ];
      passedAttrs = builtins.removeAttrs attrs specialAttrs;

      getDeps = depAttr: let
        deps = getAttrDefault depAttr pyProject.tool.poetry {};
        depAttrs = builtins.map (d: lib.toLower d) (builtins.attrNames deps);
      in
        builtins.map (dep: py.pkgs."${dep}") depAttrs;

      getInputs = attr: getAttrDefault attr attrs [];
      mkInput = attr: extraInputs: getInputs attr ++ extraInputs;

      buildSystemPkgs = poetryLib.getBuildSystemPkgs {
        inherit pyProject;
        pythonPackages = py.pkgs;
      };

    in
      py.pkgs.buildPythonApplication (
        passedAttrs // {
          pname = pyProject.tool.poetry.name;
          version = pyProject.tool.poetry.version;

          format = "pyproject";

          nativeBuildInputs = [ pkgs.yj ];
          buildInputs = mkInput "buildInputs" buildSystemPkgs;
          propagatedBuildInputs = mkInput "propagatedBuildInputs" (getDeps "dependencies") ++ ([ py.pkgs.setuptools ]);
          checkInputs = mkInput "checkInputs" (getDeps "dev-dependencies");

          passthru = {
            python = py;
          };

          postPatch = (getAttrDefault "postPatch" passedAttrs "") + ''
            # Tell poetry not to resolve the path dependencies. Any version is
            # fine !
            yj -tj < pyproject.toml | python ${./pyproject-without-path.py} > pyproject.json
            yj -jt < pyproject.json > pyproject.toml
            rm pyproject.json
          '';

          meta = meta // {
            inherit (pyProject.tool.poetry) description;
            licenses = [ pyProject.tool.poetry.license ];
          };

        }
      );

  cli = import ./cli.nix { inherit pkgs lib; };

in
{
  inherit mkPoetryEnv mkPoetryApplication defaultPoetryOverrides cli;
  mkPoetryPackage = attrs: builtins.trace "mkPoetryPackage is deprecated. Use mkPoetryApplication instead." (mkPoetryApplication attrs);
}