devenv.nix

{ config, inputs, lib, pkgs, ... }:

let
  cfg = config.languages.python.uv2nix;
in {
  options.languages.python = with lib; {
    uv2nix = {
      enable = mkEnableOption "uv2nix";
      package = mkOption {
        type = lib.types.package;
        description = "The python package to use";
      };
      packages = mkOption {
        type = with types; listOf package;
        default = [];
        description = "Packages to include on the same level as python (using this instead of packages could help with tools like pytest not working, for example you might put python tools there like pre-commit as they might mess up the environment)";
      };
      preferWheels = mkOption {
        type = types.bool;
        default = true;
        description = "whther to prefer wheels over source packages";
      };
      root = mkOption {
        type = types.path;
        description = "root of the project";
      };
      projectFiles = mkOption {
        type = with types; listOf path;
        description = "list of files that make the project, so things don't have to be recomputed on every change";
        default = [
          (cfg.root + "/pyproject.toml")
        ];
      };
      extraDependencies = mkOption {
        type = with types; attrsOf (listOf str);
        description = "list of additional python dependencies";
        default = {
          editables = [];
        };
      };
      buildSystemOverrides = mkOption {
        type = with types; attrsOf (attrsOf (listOf str));
        description = "build system overrides";
        default = {};
      };
      overrides = mkOption {
        type = types.unspecified;
        description = "advanced overrides for packages";
        default = self: super: {};
      };
      injectAppEnv = mkOption {
        type = types.bool;
        description = "include the app and its dependencies in the dev shell";
        default = true;
      };
      impureEnv = mkOption {
        type = types.bool;
        description = "use env built by uv over uv2nix";
        default = false;
      };
    };
  };
  config = let
    lockExists = builtins.pathExists (cfg.root + "/uv.lock");
    project = inputs.pyproject-nix.lib.project.loadPyproject {
      projectRoot = cfg.root;
    };
    projectName = project.pyproject.project.name;

    workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
      workspaceRoot = cfg.root;
    };

    overlay = workspace.mkPyprojectOverlay {
      sourcePreference = if cfg.preferWheels then "wheel" else "sdist";
    };
    buildSystemOverrides = final: prev: builtins.mapAttrs (
      name: spec:
      prev.${name}.overrideAttrs (old: {
        nativeBuildInputs = old.nativeBuildInputs ++ final.resolveBuildSystem spec;
      })
    ) cfg.buildSystemOverrides;

    pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
      python = cfg.package;
    }).overrideScope (
      lib.composeManyExtensions [
        inputs.pyproject-build-systems.overlays.default
        overlay
        buildSystemOverrides
        cfg.overrides
      ]
    );

    editableOverlay = workspace.mkEditablePyprojectOverlay {
      root = "$DEVENV_ROOT";
      # members = [ "hello-world" ];
    };
    editablePythonSet = pythonSet.overrideScope (
      lib.composeManyExtensions [
        editableOverlay

        (final: prev: {
          ${projectName} = prev.${projectName}.overrideAttrs (old: {
            src = lib.fileset.toSource {
              root = cfg.root;
              fileset = lib.fileset.unions cfg.projectFiles;
            };

            nativeBuildInputs = old.nativeBuildInputs ++ final.resolveBuildSystem cfg.extraDependencies;
          });
        })
      ]
    );

    app = pythonSet.mkVirtualEnv "${projectName}-env" workspace.deps.default;
    env = editablePythonSet.mkVirtualEnv "${projectName}-dev-env" workspace.deps.all;
  in lib.mkIf cfg.enable {
    cachix.pull = [ "pyproject-nix" ];

    env = {
      # don't allow uv to download different python interpreters
      UV_PYTHON_DOWNLOADS = "never";
    }
    // lib.optionalAttrs cfg.impureEnv (
      {
        # make uv use python provided by nix
        UV_PYTHON = config.languages.python.package;
      } // lib.optionalAttrs pkgs.stdenv.isLinux {
        LD_LIBRARY_PATH = lib.makeLibraryPath pkgs.pythonManylinuxPackages.manylinux1;
      }
    )
    // lib.optionalAttrs (!cfg.impureEnv) {
      UV_NO_SYNC = "1";
      UV_PYTHON = if lockExists then "${config.outputs.python.env}/bin/python" else "${config.languages.python.package}/bin/python";
    };
    languages.python.uv.package = lib.mkDefault inputs.uv2nix.packages.${pkgs.stdenv.targetPlatform.system}.uv-bin;
    outputs.python = {
      inherit app env editableOverlay editablePythonSet config;
    };
    packages = [
      config.languages.python.uv.package
    ] ++ lib.optional (lockExists && cfg.injectAppEnv) config.outputs.python.env;
    unsetEnvVars = ["PYTHONPATH"];
  };
}