Fix mypy errors for python 3.9

Author: AlePouroullis

src/humanloop/client.py

@@ -395,7 +395,7 @@ def agent():
             attributes=attributes,
         )
 
-    def pull(self, path: str | None = None, environment: str | None = None) -> Tuple[List[str], List[str]]:
+    def pull(self, path: Optional[str] = None, environment: Optional[str] = None) -> Tuple[List[str], List[str]]:
         """Pull Prompt and Agent files from Humanloop to local filesystem.
 
         This method will:

src/humanloop/sync/sync_client.py

@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import List, Tuple, TYPE_CHECKING
+from typing import List, Optional, Tuple, TYPE_CHECKING, Union
 from functools import lru_cache
 import typing
 import time
@@ -203,7 +203,7 @@ def _save_serialized_file(
             logger.error(f"Failed to write {file_type} {file_path} to disk: {str(e)}")
             raise
 
-    def _pull_file(self, path: str, environment: str | None = None) -> bool:
+    def _pull_file(self, path: str, environment: Optional[str] = None) -> bool:
         """Pull a specific file from Humanloop to local filesystem.
 
         Returns:
@@ -236,8 +236,8 @@ def _pull_file(self, path: str, environment: str | None = None) -> bool:
 
     def _pull_directory(
         self,
-        path: str | None = None,
-        environment: str | None = None,
+        path: Optional[str] = None,
+        environment: Optional[str] = None,
     ) -> Tuple[List[str], List[str]]:
         """Sync Prompt and Agent files from Humanloop to local filesystem.
 
@@ -316,7 +316,7 @@ def _pull_directory(
 
         return successful_files, failed_files
 
-    def pull(self, path: str | None = None, environment: str | None = None) -> Tuple[List[str], List[str]]:
+    def pull(self, path: Optional[str] = None, environment: Optional[str] = None) -> Tuple[List[str], List[str]]:
         """Pull files from Humanloop to local filesystem.
 
         If the path ends with .prompt or .agent, pulls that specific file.
@@ -343,7 +343,9 @@ def pull(self, path: str | None = None, environment: str | None = None) -> Tuple
         )
 
         try:
-            if normalized_path is None or path is None: # path being None means normalized_path is None, but we check both for improved type safety
+            if (
+                normalized_path is None or path is None
+            ):  # path being None means normalized_path is None, but we check both for improved type safety
                 # Pull all files from the root
                 logger.debug("Pulling all files from root")
                 successful_files, failed_files = self._pull_directory(

tests/custom/integration/conftest.py

@@ -1,10 +1,10 @@
-import io
 import os
 import time
+import typing
 import uuid
-from contextlib import contextmanager, redirect_stdout
+from collections.abc import Generator
 from dataclasses import dataclass
-from typing import ContextManager, Generator, List, TextIO, Union
+from typing import Union
 
 import dotenv
 import pytest
@@ -21,17 +21,6 @@ class ResourceIdentifiers:
     file_path: str
 
 
-@pytest.fixture()
-def capture_stdout() -> ContextManager[TextIO]:
-    @contextmanager
-    def _context_manager():
-        f = io.StringIO()
-        with redirect_stdout(f):
-            yield f
-
-    return _context_manager  # type: ignore [return-value]
-
-
 @pytest.fixture(scope="session")
 def openai_key() -> str:
     dotenv.load_dotenv()
@@ -44,26 +33,26 @@ def openai_key() -> str:
 def sdk_test_dir(get_humanloop_client: GetHumanloopClientFn) -> Generator[str, None, None]:
     humanloop_client = get_humanloop_client()
 
+    def _get_subclient(file_type: str):
+        try:
+            return {
+                "agent": humanloop_client.agents,
+                "prompt": humanloop_client.prompts,
+                "dataset": humanloop_client.datasets,
+                "evaluator": humanloop_client.evaluators,
+                "flow": humanloop_client.flows,
+                "tool": humanloop_client.tools,
+            }[file_type]
+        except KeyError:
+            raise NotImplementedError(f"Unknown file type: {file_type}")
+
     def cleanup_directory(directory_id: str):
         directory_response = humanloop_client.directories.get(id=directory_id)
         for subdirectory in directory_response.subdirectories:
             cleanup_directory(subdirectory.id)
         for file in directory_response.files:
-            match file.type:
-                case "agent":
-                    humanloop_client.agents.delete(id=file.id)
-                case "prompt":
-                    humanloop_client.prompts.delete(id=file.id)
-                case "dataset":
-                    humanloop_client.datasets.delete(id=file.id)
-                case "evaluator":
-                    humanloop_client.evaluators.delete(id=file.id)
-                case "flow":
-                    humanloop_client.flows.delete(id=file.id)
-                case "tool":
-                    humanloop_client.tools.delete(id=file.id)
-                case _:
-                    raise ValueError(f"Unknown file type: {file.type}")
+            subclient = _get_subclient(typing.cast(str, file.type))
+            subclient.delete(id=file.id)
         humanloop_client.directories.delete(id=directory_response.id)
 
     path = f"SDK_INTEGRATION_TEST_{uuid.uuid4()}"
@@ -211,7 +200,7 @@ def syncable_files_fixture(
     sdk_test_dir: str,
 ) -> Generator[list[SyncableFile], None, None]:
     """Creates a predefined structure of files in Humanloop for testing sync."""
-    files: List[SyncableFile] = [
+    files: list[SyncableFile] = [
         SyncableFile(
             path="prompts/gpt-4",
             type="prompt",

tests/custom/integration/test_sync.py

@@ -1,3 +1,4 @@
+import typing
 from pathlib import Path
 from typing import List, Union
 
@@ -80,15 +81,17 @@ def test_overload_with_local_files(
     # WHEN calling with an invalid path
     # THEN it should raise HumanloopRuntimeError
     with pytest.raises(HumanloopRuntimeError):
-        sub_client: Union[PromptsClient, AgentsClient]
-        match test_file.type:
-            case "prompt":
-                sub_client = humanloop_client.prompts
-            case "agent":
-                sub_client = humanloop_client.agents
-            case _:
-                raise ValueError(f"Invalid file type: {test_file.type}")
-        sub_client.call(path="invalid/path")
+        try:
+            sub_client: Union[PromptsClient, AgentsClient] = typing.cast(
+                Union[PromptsClient, AgentsClient],
+                {
+                    "prompt": humanloop_client.prompts,
+                    "agent": humanloop_client.agents,
+                }[test_file.type],
+            )
+            sub_client.call(path="invalid/path")
+        except KeyError:
+            raise NotImplementedError(f"Unknown file type: {test_file.type}")
 
 
 def test_overload_log_with_local_files(