flake-python/flake.nix

{
  description = "Customizable Python development environment";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, flake-utils }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        # Function to create a Python environment with a given version
        mkPython = {
          pythonVersion ? "python313",  # Default Python version
          extraPackages ? [], # Additional packages to include
          extraSystemPackages ? [] # Additional system packages
        }:
        let
          pkgs = import nixpkgs { inherit system; };

          # Select Python version based on the parameter
          python = pkgs.${pythonVersion};
          
          # Define package overrides for specific Python versions
          packageOverrides = self: super: {
            # Fix Sphinx for Python 3.10
            sphinx = if pythonVersion == "python310" then
              super.sphinx.overridePythonAttrs (old: { 
                version = "latest";  # Use a version known to work with Python 3.10
              })
            else
              super.sphinx;
            
            # Add more package overrides as needed
          };
          
          # Create Python with package overrides
          pythonWithOverrides = python.override {
            packageOverrides = packageOverrides;
          };

          # Base Python packages that are always included
          basePackages = ps: with ps; [
            # Data science
            numpy
            pandas
            matplotlib

            # Development tools
            black
            pytest
            ipython

            # Add more default packages as needed
          ];

          # Combine base packages with extra packages
          allPackages = ps: (basePackages ps) ++ (extraPackages ps);

          # Create the Python environment using the Python with overrides
          pythonEnv = pythonWithOverrides.withPackages allPackages;
        in
          {
            # The Python environment
            package = pythonEnv;

            # Make the chosen Python version accessible
            python = pythonWithOverrides;
            inherit pythonVersion;

            # Development shell with the Python environment
            devShell = pkgs.mkShell {
              buildInputs = [
                pythonEnv
              ] ++ extraSystemPackages;

              shellHook = ''
                echo "Python $(${pythonWithOverrides}/bin/python --version)"
                echo "Custom Python environment activated!"
              '';
            };
          };

        # Default Python configuration
        defaultPython = mkPython {};
      in
      {
        # Expose the function to create custom Python environments
        lib = {
          inherit mkPython;
        };

        # Default packages using default Python version
        packages = {
          default = defaultPython.package;
        };

        # Default development shell
        devShells.default = defaultPython.devShell;

        # NixOS module for system-wide integration
        nixosModules.default = { config, lib, pkgs, ... }:
          with lib;
          let
            cfg = config.services.customPython;
          in {
            options.services.customPython = {
              enable = mkEnableOption "Enable the custom Python environment";

              pythonVersion = mkOption {
                type = types.str;
                default = "python311";
                description = "Python version to use (e.g., python39, python310, python311)";
              };

              extraPackages = mkOption {
                type = types.functionTo (types.listOf types.anything);
                default = _: [];
                description = "Extra Python packages to include";
              };

              extraSystemPackages = mkOption {
                type = types.listOf types.package;
                default = [];
                description = "Extra system packages to include";
              };

              includeInSystemPackages = mkOption {
                type = types.bool;
                default = true;
                description = "Whether to include the Python environment in system packages";
              };
            };

            config = mkIf cfg.enable {
              environment.systemPackages = mkIf cfg.includeInSystemPackages [
                (mkPython {
                  pythonVersion = cfg.pythonVersion;
                  extraPackages = cfg.extraPackages;
                  extraSystemPackages = cfg.extraSystemPackages;
                }).package
              ];
            };
          };
      }
    );
}