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};

          # Determine if Python version is 3.11 or higher
          isPython311OrHigher = 
            let
              versionStr = builtins.substring 6 10 pythonVersion;
              versionNum = builtins.fromJSON (builtins.substring 0 1 versionStr + "." + builtins.substring 1 2 versionStr);
            in
              versionNum >= 3.11;

          # Custom function to filter out sphinx when Python version is < 3.11
          filterSphinxForOldPython = packages:
            builtins.filter (p: p.pname or "" != "sphinx") packages;

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

            # Development tools
            black
            pytest
            ipython

            # Include sphinx (will be filtered out later if needed)
            sphinx

            # Add more default packages as needed
          ];

          # Combine base packages with extra packages and filter if needed
          allPackages = ps: 
            let
              combinedPackages = (basePackages ps) ++ (extraPackages ps);
            in
              filterSphinxForOldPython combinedPackages;

          # Create the Python environment with filtered packages
          pythonEnv = python.withPackages allPackages;
        in
          {
            # The Python environment
            package = pythonEnv;

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

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

              shellHook = ''
                echo "Python $(${python}/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 = "python313";
                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
              ];
            };
          };
      }
    );
}