From 442b561d3c4e149afb362ce0ddaeba95a618d33a Mon Sep 17 00:00:00 2001
From: mseitzer <16725193+mseitzer@users.noreply.github.com>
Date: Sun, 17 Mar 2024 14:00:35 +0100
Subject: [PATCH] Run nox sessions through poetry versions

---
 .github/workflows/tests_full.yaml    | 24 ++++++++++++++++------
 .github/workflows/tests_reduced.yaml | 16 +++++++++++++--
 noxfile.py                           | 30 ++++++++++++++++++----------
 3 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/tests_full.yaml b/.github/workflows/tests_full.yaml
index cf1949b..29d11ce 100644
--- a/.github/workflows/tests_full.yaml
+++ b/.github/workflows/tests_full.yaml
@@ -11,7 +11,7 @@ concurrency:
   
 jobs:
   tests:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
@@ -22,20 +22,32 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+    - name: Install Poetry
+      run: |
+        pip install poetry==8.2.0
+        pip install poetry-plugin-export==1.7.0
     - name: Install Nox
-      run: pip install nox==2024.03.02
+      run: |
+        pip install nox==2024.03.02
+        pip install nox-poetry==1.0.3
     - name: Run tests
       run: nox --non-interactive --error-on-missing-interpreter --session "tests-${{ matrix.python-version }}"
   
   lint:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.9
+      - name: Set up Python 3.11
         uses: actions/setup-python@v5
         with:
-          python-version: 3.9
+          python-version: 3.11
+      - name: Install Poetry
+        run: |
+          pip install poetry==8.2.0
+          pip install poetry-plugin-export==1.7.0
       - name: Install Nox
-        run: pip install nox==2024.03.02
+        run: |
+          pip install nox==2024.03.02
+          pip install nox-poetry==1.0.3
       - name: Lint
         run: nox --non-interactive --error-on-missing-interpreter --session "lint"
diff --git a/.github/workflows/tests_reduced.yaml b/.github/workflows/tests_reduced.yaml
index a645b5a..9fc9f36 100644
--- a/.github/workflows/tests_reduced.yaml
+++ b/.github/workflows/tests_reduced.yaml
@@ -21,8 +21,14 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+    - name: Install Poetry
+      run: |
+          pip install poetry==8.2.0
+          pip install poetry-plugin-export==1.7.0
     - name: Install Nox
-      run: pip install nox==2024.03.02
+      run: |
+        pip install nox==2024.03.02
+        pip install nox-poetry==1.0.3
     - name: Run tests
       run: nox --non-interactive --error-on-missing-interpreter --session "tests-${{ matrix.python-version }}"
   
@@ -34,7 +40,13 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: 3.9
+      - name: Install Poetry
+        run: |
+          pip install poetry==8.2.0
+          pip install poetry-plugin-export==1.7.0
       - name: Install Nox
-        run: pip install nox==2024.03.02
+        run: |
+          pip install nox==2024.03.02
+          pip install nox-poetry==1.0.3
       - name: Lint
         run: nox --non-interactive --error-on-missing-interpreter --session "lint"
diff --git a/noxfile.py b/noxfile.py
index 8dc5c82..d236f64 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -1,29 +1,39 @@
+import sys
+from textwrap import dedent
+
 import nox
 
+
+try:
+    from nox_poetry import session
+except ImportError:
+    message = f"""\
+    Nox failed to import the 'nox-poetry' package.
+
+    Please install it using the following command:
+
+    {sys.executable} -m pip install nox-poetry"""
+    raise SystemExit(dedent(message)) from None
+
+
 LOCATIONS = ("src/", "tests/", "noxfile.py", "setup.py")
 
 
-@nox.session
+@session
 def lint(session):
     session.install("flake8")
     session.install("flake8-bugbear")
     session.install("flake8-isort")
-    session.install("black==24.3.0")
+    session.install("black")
 
     args = session.posargs or LOCATIONS
     session.run("flake8", *args)
     session.run("black", "--check", "--diff", *args)
 
 
-@nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"])
+@session(python=["3.8", "3.9", "3.10", "3.11", "3.12"])
 def tests(session):
-    session.install(
-        "torch==2.2.1",
-        "torchvision",
-        "--index-url",
-        "https://download.pytorch.org/whl/cpu",
-    )
-    session.install(".")
+    session.install(".", '--extra-index-url', 'https://download.pytorch.org/whl/cpu')
     session.install("pytest")
     session.install("pytest-mock")
     session.run("pytest", *session.posargs)