refactor: run mypy on unit tests (#5132)

Co-authored-by: Roger Yang <[email protected]>
Arize-ai · Oct 22, 2024 · 035e394 · 035e394
1 parent 083ef42
commit 035e394
Show file tree

Hide file tree

Showing 80 changed files with 514 additions and 345 deletions.
diff --git a/.github/workflows/playwright.yaml b/.github/workflows/playwright.yaml
@@ -53,7 +53,7 @@ jobs:
           python-version: 3.12
       - uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: "pyproject.toml"
           github-token: ${{ secrets.GITHUB_TOKEN }}

diff --git a/.github/workflows/python-CI.yml b/.github/workflows/python-CI.yml
@@ -98,7 +98,7 @@ jobs:
           python-version: ${{ matrix.py }}
       - uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
       - run: uv tool install tox --with tox-uv
       - run: tox run -e phoenix_evals -- -ra -x
 
@@ -122,7 +122,7 @@ jobs:
           python-version: ${{ matrix.py }}
       - uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
       - run: uv tool install tox --with tox-uv
       - run: tox run -e phoenix_otel -- -ra -x
 
@@ -145,7 +145,7 @@ jobs:
   #      - name: Set up `uv`
   #        uses: astral-sh/setup-uv@v3
   #        with:
-  #          version: 0.4.22
+  #          version: 0.4.25
   #          enable-cache: true
   #          cache-dependency-glob: |
   #            pyproject.toml
@@ -176,7 +176,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -207,7 +207,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -238,7 +238,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -267,7 +267,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -295,14 +295,15 @@ jobs:
           sparse-checkout: |
             requirements/
             src/phoenix/
+            tests/
       - name: Set up Python ${{ matrix.py }}
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.py }}
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -314,6 +315,42 @@ jobs:
       - name: Check types
         run: tox run -e type_check
 
+  type-check-unit-tests:
+    name: Type Check Unit Tests
+    runs-on: ${{ matrix.os }}
+    needs: changes
+    if: ${{ needs.changes.outputs.phoenix == 'true' }}
+    strategy:
+      matrix:
+        py: [3.9, 3.12]
+        os: [ubuntu-latest]
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          sparse-checkout: |
+            requirements/
+            src/phoenix/
+            tests/
+      - name: Set up Python ${{ matrix.py }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.py }}
+      - name: Set up `uv`
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: 0.4.25
+          enable-cache: true
+          cache-dependency-glob: |
+            pyproject.toml
+            requirements/ci.txt
+            requirements/unit-tests.txt
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Set up `tox` with `tox-uv`
+        run: uv tool install tox --with tox-uv
+      - name: Check types on unit tests
+        run: tox run -e type_check_unit_tests
+
   unit-tests:
     name: Unit Tests
     runs-on: ${{ matrix.os }}
@@ -339,7 +376,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -382,7 +419,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml
@@ -440,7 +477,7 @@ jobs:
       - name: Set up `uv`
         uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml

diff --git a/.github/workflows/python-cron.yml b/.github/workflows/python-cron.yml
@@ -31,7 +31,7 @@ jobs:
           python-version: 3.9
       - uses: astral-sh/setup-uv@v3
         with:
-          version: 0.4.22
+          version: 0.4.25
           enable-cache: true
           cache-dependency-glob: |
             pyproject.toml

diff --git a/pyproject.toml b/pyproject.toml
@@ -77,7 +77,7 @@ dev = [
   "jupyter",
   "nbqa",
   "ruff==0.6.9",
-  "mypy==1.11.2",
+  "mypy==1.12.1",
   "pandas>=1.0",
   "tabulate",  # used by DataFrame.to_markdown()
   "types-tabulate",
@@ -215,7 +215,7 @@ dependencies = [
 dependencies = [
   "grpcio",
   "litellm>=1.0.3",
-  "mypy==1.11.2",
+  "mypy==1.12.1",
   "openai>=1.0.0",
   "opentelemetry-exporter-otlp",
   "opentelemetry-instrumentation-fastapi",

diff --git a/requirements/ci.txt b/requirements/ci.txt
@@ -1,4 +1,4 @@
-mypy==1.11.2
+mypy==1.12.1
 pytest==8.3.3
 pytest-xdist==3.6.1
 pytest-asyncio==0.24.0

diff --git a/requirements/integration-tests.txt b/requirements/integration-tests.txt
@@ -1,11 +1,13 @@
 -r ci.txt
+asyncpg
 beautifulsoup4
 faker
 httpx
 openinference-semantic-conventions
 opentelemetry-sdk
 portpicker
 psutil
+psycopg[binary,pool]
 pyjwt
 pytest-randomly
 pytest-smtpd

diff --git a/requirements/type-check.txt b/requirements/type-check.txt
@@ -1,5 +1,6 @@
 -r ci.txt
 anthropic
+asyncpg
 grpcio
 litellm>=1.0.3
 openai>=1.0.0
@@ -13,6 +14,7 @@ opentelemetry-semantic-conventions
 pandas-stubs==2.0.3.230814
 pandas>=1.0
 prometheus_client
+psycopg[binary,pool]
 py-grpc-prometheus
 pypistats  # this is needed to type-check third-party packages
 requests  # this is needed to type-check third-party packages

diff --git a/requirements/unit-tests.txt b/requirements/unit-tests.txt
@@ -10,12 +10,16 @@ litellm>=1.0.3
 nest-asyncio # for executor testing
 numpy
 openai>=1.0.0
+pandas-stubs==2.0.3.230814
+pandas>=1.0
 protobuf==3.20.2  # version minimum (for tests)
 psycopg[binary,pool]
+pyarrow-stubs
 pytest-postgresql
 responses
 respx # For OpenAI testing
 tenacity
 tiktoken
+types-pytz
 typing-extensions==4.7.0
 vcrpy
diff --git a/src/phoenix/core/model_schema.py b/src/phoenix/core/model_schema.py
@@ -1175,7 +1175,7 @@ class Schema(SchemaSpec):
     prediction_score: Optional[str] = None
     actual_label: Optional[str] = None
     actual_score: Optional[str] = None
-    prompt: Optional[RetrievalEmbedding] = None
+    prompt: Optional[Embedding] = None
     response: Optional[Union[str, Embedding]] = None
     features: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)
     tags: Iterable[Union[str, CompositeDimensionSpec]] = field(default_factory=list)

diff --git a/src/phoenix/pointcloud/pointcloud.py b/src/phoenix/pointcloud/pointcloud.py
@@ -1,9 +1,8 @@
 from dataclasses import dataclass
-from typing import Dict, List, Mapping, Protocol, Set, Tuple
+from typing import Dict, Hashable, List, Mapping, Protocol, Set, Tuple, TypeVar
 
 import numpy as np
 import numpy.typing as npt
-from strawberry import ID
 from typing_extensions import TypeAlias
 
 from phoenix.pointcloud.clustering import RawCluster
@@ -12,6 +11,8 @@
 Matrix: TypeAlias = npt.NDArray[np.float64]
 RowIndex: TypeAlias = int
 
+_IdType = TypeVar("_IdType", bound=Hashable)
+
 
 class DimensionalityReducer(Protocol):
     def project(self, mat: Matrix, n_components: int) -> Matrix: ...
@@ -28,9 +29,9 @@ class PointCloud:
 
     def generate(
         self,
-        data: Mapping[ID, Vector],
+        data: Mapping[_IdType, Vector],
         n_components: int = 3,
-    ) -> Tuple[Dict[ID, Vector], Dict[str, Set[ID]]]:
+    ) -> Tuple[Dict[_IdType, Vector], Dict[str, Set[_IdType]]]:
         """
         Given a set of vectors, projects them onto lower dimensions, and
         finds clusters among the projections.

diff --git a/src/phoenix/server/api/types/Project.py b/src/phoenix/server/api/types/Project.py
@@ -223,18 +223,13 @@ async def spans(
             span_records = await session.execute(stmt)
             async for span_record in islice(span_records, first):
                 span = span_record[0]
-                sort_column_value = span_record[1] if len(span_record) > 1 else None
-                cursor = Cursor(
-                    rowid=span.id,
-                    sort_column=(
-                        CursorSortColumn(
-                            type=sort_config.column_data_type,
-                            value=sort_column_value,
-                        )
-                        if sort_config
-                        else None
-                    ),
-                )
+                cursor = Cursor(rowid=span.id)
+                if sort_config:
+                    assert len(span_record) > 1
+                    cursor.sort_column = CursorSortColumn(
+                        type=sort_config.column_data_type,
+                        value=span_record[1],
+                    )
                 cursors_and_nodes.append((cursor, to_gql_span(span)))
             has_next_page = True
             try:

diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,4 +1,4 @@
-from _pytest.config import Parser
+from _pytest.config.argparsing import Parser
 
 
 def pytest_addoption(parser: Parser) -> None:

diff --git a/tests/integration/mypy.ini b/tests/integration/mypy.ini
@@ -1,3 +1,2 @@
 [mypy]
 strict = true
-exclude = (^evals|^notebooks)
diff --git a/tests/mypy.ini b/tests/mypy.ini
@@ -0,0 +1,16 @@
+[mypy]
+strict = true
+exclude = (?x)(
+          ^integration/$ # integration tests
+          | ^unit/datasets/test_experiments\.py$ # TODO: fix type errors
+          | ^unit/server/api/routers/v1/test_annotations\.py$ # TODO: fix type errors
+          )
+
+[mypy-nest_asyncio.*]
+ignore_missing_imports = true
+
+[mypy-sklearn.*]
+ignore_missing_imports = true
+
+[mypy-vcr.*]
+ignore_missing_imports = true
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py