name: Python

on:
  pull_request:

jobs:
  build_and_test:
    name: Check everything builds & tests
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: [ubuntu-latest, windows-latest]
        # Lowest and highest, no version specified so that
        # new releases get automatically tested against
        version: [
          {torch: torch==2.4, python: "3.10", arch: "x64", numpy: numpy==1.26.4},
          {torch: torch, python: "3.13", arch: "x64", numpy: numpy}
        ]
        include:
          - os: ubuntu-latest
            version:
              torch: torch
              python: "3.14"
              numpy: numpy
          - os: ubuntu-latest
            version:
              torch: torch
              python: "3.14"
              numpy: numpy
              arch: "x64-freethreaded"
          - os: macos-15-intel
            version:
              torch: torch
              numpy: "numpy==1.26"
              python: "3.10"
              arch: "x64"
          - os: macos-latest
            version:
              torch: torch
              python: "3.13"
              numpy: numpy
              arch: "arm64"
          - os: windows-11-arm
            version:
              torch: torch
              python: "3.13"
              numpy: numpy
              arch: "arm64"
    defaults:
      run:
        working-directory: ./bindings/python
    steps:
      - name: Checkout repository
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2

      - name: Install Rust
        uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8  # stable
        with:
          components: rustfmt, clippy

      - name: Cargo install audit
        run: cargo install cargo-audit

      - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32  # v2
        with:
          workspaces: "bindings/python"

      - name: Install Python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6
        with:
          python-version: ${{ matrix.version.python }}
          architecture: ${{ matrix.version.arch }}

      - name: Lint with RustFmt
        run: cargo fmt -- --check

      - name: Lint with Clippy
        run: cargo clippy --all-targets --all-features -- -D warnings

      - name: Run Audit
        run: cargo audit -D warnings

      - name: Install (torch)
        if: matrix.os != 'windows-11-arm'
        run: |
          pip install ${{ matrix.version.numpy }}
          pip install ${{ matrix.version.torch }}
        shell: bash

      - name: Install (torch windows arm64)
        if: matrix.os == 'windows-11-arm'
        run: |
          pip install ${{ matrix.version.numpy }}
          pip install ${{ matrix.version.torch }} --index-url https://download.pytorch.org/whl/cpu
        shell: bash

      - name: Install (hdf5 non windows)
        if: matrix.os == 'ubuntu-latest'
        run: |
          sudo apt-get update
          sudo apt-get install libhdf5-dev

      - name: Install (tensorflow)
        if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
        run: |
          pip install ".[${{ matrix.version.python == '3.14' && 'tf-nightly' || 'tensorflow' }}]"
          # XXX: I assume this is still the case for 3.10
          # Force reinstall of numpy, tensorflow uses numpy 2 even on 3.9
          pip install ${{ matrix.version.numpy }}
        shell: bash

      - name: Install (jax, flax)
        if: runner.os != 'Windows'
        run:
          pip install .[jax]
        shell: bash

      - name: Install (mlx)
        if: matrix.os == 'macos-latest'
        run: |
          pip install .[mlx]
        shell: bash

      - name: Check style
        run: |
          pip install .[quality]
          ruff format --check .

      - name: Run tests
        if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
        run: |
          cargo test
          pip install ".[testing]"
          pytest -sv tests/

      - name: Run tests (Windows arm64)
        if: matrix.os == 'windows-11-arm'
        run: |
          cargo test
          pip install ".[testing]"
          pytest -sv tests/ --ignore=tests/test_tf_comparison.py

      - name: Run tests (freethreaded)
        if: matrix.version.arch == 'x64-freethreaded'
        run: |
          cargo test
          pip install --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple ".[testing]"
          # importing msgpack._cmsgpack re-enables the GIL, so force-disable
          # until msgpack declares support
          PYTHON_GIL=0 pytest -sv tests/ --ignore=tests/test_tf_comparison.py

  test_s390x_big_endian:
    runs-on: ubuntu-latest
    permissions:
      packages: write
    name: Test bigendian - S390X
    steps:
      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
      - name: Set up QEMU
        uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130  # v3
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f  # v3
      - name: Can push to GHCR?
        id: canpush
        shell: bash
        run: |
          echo "value=${{ github.event.pull_request.head.repo.fork == false }}" >> "$GITHUB_OUTPUT"
      - name: Docker meta
        id: meta
        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051  # v5
        with:
          # list of Docker images to use as base name for tags
          images: |
            ghcr.io/safetensors/safetensors/s390x
          # generate Docker tags based on the following events/attributes
          tags: |
            type=schedule
            type=ref,event=branch
            type=ref,event=pr
            type=semver,pattern={{version}}
            type=semver,pattern={{major}}.{{minor}}
            type=semver,pattern={{major}}
            type=sha
      - name: Login to Registry
        if: steps.canpush.outputs.value == 'true'
        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9  # v3
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}
      - name: Test big endian
        uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8  # v6
        with:
          platforms: linux/s390x
          file: Dockerfile.s390x.test
          tags: ${{ steps.meta.outputs.tags }}
          labels: ${{ steps.meta.outputs.labels }}
          cache-from: ${{ steps.canpush.outputs.value == 'true' && 'type=registry,ref=ghcr.io/safetensors/safetensors/s390x:cache,mode=max' || 'type=gha' }}
          cache-to: ${{ steps.canpush.outputs.value == 'true' && 'type=registry,ref=ghcr.io/safetensors/safetensors/s390x:cache,mode=max' || 'type=gha' }}
          push: ${{ steps.canpush.outputs.value == 'true' }}
