Add pyupgrade --py39-plus and autoflake to pre-commit config (#1329)

* Update `pyupgrade` in `pre-commit` to `--py39-plus` to enforce PEP 585

Also add `autoflake` to remove unused imports after switching to PEP 585.

* pre-commit

* pre-commit

---------

Co-authored-by: David L. Qiu <[email protected]>

Author: rominf

.pre-commit-config.yaml

@@ -17,6 +17,11 @@ repos:
       - id: check-builtin-literals
       - id: trailing-whitespace
 
+  - repo: https://github.com/PyCQA/autoflake
+    rev: v2.3.1
+    hooks:
+      - id: autoflake
+
   - repo: https://github.com/psf/black
     rev: 25.1.0
     hooks:
@@ -33,7 +38,7 @@ repos:
     rev: v3.19.1
     hooks:
       - id: pyupgrade
-        args: [--py37-plus]
+        args: [--py39-plus]
 
   - repo: https://github.com/pycqa/flake8
     rev: 7.1.1

packages/jupyter-ai-magics/jupyter_ai_magics/completion_utils.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 from .models.completion import InlineCompletionRequest
 
 
@@ -9,7 +7,7 @@ def token_from_request(request: InlineCompletionRequest, suggestion: int):
     return f"t{request.number}s{suggestion}"
 
 
-def template_inputs_from_request(request: InlineCompletionRequest) -> Dict:
+def template_inputs_from_request(request: InlineCompletionRequest) -> dict:
     suffix = request.suffix.strip()
     filename = request.path.split("/")[-1] if request.path else "untitled"
 

packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, List, Optional
+from typing import ClassVar, Optional
 
 from jupyter_ai_magics.providers import (
     AuthStrategy,
@@ -27,7 +27,7 @@ class BaseEmbeddingsProvider(BaseModel):
     name: ClassVar[str] = ...
     """User-facing name of this provider."""
 
-    models: ClassVar[List[str]] = ...
+    models: ClassVar[list[str]] = ...
     """List of supported models by their IDs. For registry providers, this will
     be just ["*"]."""
 
@@ -38,7 +38,7 @@ class BaseEmbeddingsProvider(BaseModel):
     model_id_key: ClassVar[str] = ...
     """Kwarg expected by the upstream LangChain provider."""
 
-    pypi_package_deps: ClassVar[List[str]] = []
+    pypi_package_deps: ClassVar[list[str]] = []
     """List of PyPi package dependencies."""
 
     auth_strategy: ClassVar[AuthStrategy] = None
@@ -50,7 +50,7 @@ class BaseEmbeddingsProvider(BaseModel):
     registry: ClassVar[bool] = False
     """Whether this provider is a registry provider."""
 
-    fields: ClassVar[List[Field]] = []
+    fields: ClassVar[list[Field]] = []
     """Fields expected by this provider in its constructor. Each `Field` `f`
     should be passed as a keyword argument, keyed by `f.key`."""
 

packages/jupyter-ai-magics/jupyter_ai_magics/exception.py

@@ -1,5 +1,3 @@
-import traceback
-
 from IPython.core.magic import register_line_magic
 from IPython.core.ultratb import ListTB
 

packages/jupyter-ai-magics/jupyter_ai_magics/models/completion.py

@@ -1,4 +1,4 @@
-from typing import List, Literal, Optional
+from typing import Literal, Optional
 
 from pydantic import BaseModel
 
@@ -50,7 +50,7 @@ class CompletionError(BaseModel):
 class InlineCompletionList(BaseModel):
     """Reflection of JupyterLab's `IInlineCompletionList`."""
 
-    items: List[InlineCompletionItem]
+    items: list[InlineCompletionItem]
 
 
 class InlineCompletionReply(BaseModel):

packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/aws.py

@@ -1,6 +1,7 @@
 import copy
 import json
-from typing import Any, Coroutine, Dict
+from collections.abc import Coroutine
+from typing import Any
 
 from jsonpath_ng import parse
 from langchain_aws import BedrockEmbeddings, BedrockLLM, ChatBedrock, SagemakerEndpoint
@@ -155,7 +156,7 @@ def __init__(self, request_schema, response_path):
         self.response_path = response_path
         self.response_parser = parse(response_path)
 
-    def replace_values(self, old_val, new_val, d: Dict[str, Any]):
+    def replace_values(self, old_val, new_val, d: dict[str, Any]):
         """Replaces values of a dictionary recursively."""
         for key, val in d.items():
             if val == old_val:
@@ -165,7 +166,7 @@ def replace_values(self, old_val, new_val, d: Dict[str, Any]):
 
         return d
 
-    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
+    def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
         request_obj = copy.deepcopy(self.request_schema)
         self.replace_values("<prompt>", prompt, request_obj)
         request = json.dumps(request_obj).encode("utf-8")

packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openrouter.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 from jupyter_ai_magics import BaseProvider
 from jupyter_ai_magics.providers import EnvAuthStrategy, TextField
 from langchain_core.utils import get_from_dict_or_env
@@ -8,7 +6,7 @@
 
 class ChatOpenRouter(ChatOpenAI):
     @property
-    def lc_secrets(self) -> Dict[str, str]:
+    def lc_secrets(self) -> dict[str, str]:
         return {"openai_api_key": "OPENROUTER_API_KEY"}
 
 

packages/jupyter-ai-magics/jupyter_ai_magics/providers.py

@@ -3,15 +3,12 @@
 import functools
 import io
 import json
+from collections.abc import AsyncIterator, Coroutine
 from concurrent.futures import ThreadPoolExecutor
 from types import MappingProxyType
 from typing import (
     Any,
-    AsyncIterator,
     ClassVar,
-    Coroutine,
-    Dict,
-    List,
     Literal,
     Optional,
     Union,
@@ -127,7 +124,7 @@ class MultiEnvAuthStrategy(BaseModel):
     """Require multiple auth tokens via multiple environment variables."""
 
     type: Literal["multienv"] = "multienv"
-    names: List[str]
+    names: list[str]
 
 
 class AwsAuthStrategy(BaseModel):
@@ -183,7 +180,7 @@ class BaseProvider(BaseModel):
     name: ClassVar[str] = ...
     """User-facing name of this provider."""
 
-    models: ClassVar[List[str]] = ...
+    models: ClassVar[list[str]] = ...
     """List of supported models by their IDs. For registry providers, this will
     be just ["*"]."""
 
@@ -207,7 +204,7 @@ class BaseProvider(BaseModel):
     If unset, the label shown in the UI defaults to "Model ID".
     """
 
-    pypi_package_deps: ClassVar[List[str]] = []
+    pypi_package_deps: ClassVar[list[str]] = []
     """List of PyPi package dependencies."""
 
     auth_strategy: ClassVar[AuthStrategy] = None
@@ -217,7 +214,7 @@ class BaseProvider(BaseModel):
     registry: ClassVar[bool] = False
     """Whether this provider is a registry provider."""
 
-    fields: ClassVar[List[Field]] = []
+    fields: ClassVar[list[Field]] = []
     """User inputs expected by this provider when initializing it. Each `Field` `f`
     should be passed in the constructor as a keyword argument, keyed by `f.key`."""
 
@@ -266,7 +263,7 @@ def completion_models(self):
     # instance attrs
     #
     model_id: str
-    prompt_templates: Dict[str, PromptTemplate]
+    prompt_templates: dict[str, PromptTemplate]
     """Prompt templates for each output type. Can be overridden with
     `update_prompt_template`. The function `prompt_template`, in the base class,
     refers to this."""
@@ -599,7 +596,7 @@ class HfHubProvider(BaseProvider, HuggingFaceEndpoint):
 
     # Handle text and image outputs
     def _call(
-        self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any
+        self, prompt: str, stop: Optional[list[str]] = None, **kwargs: Any
     ) -> str:
         """Call out to Hugging Face Hub's inference endpoint.
 

packages/jupyter-ai-magics/jupyter_ai_magics/tests/test_magics.py

@@ -119,5 +119,5 @@ def test_reset(ip):
     ip.extension_manager.load_extension("jupyter_ai_magics")
     ai_magics = ip.magics_manager.registry["AiMagics"]
     ai_magics.transcript = [AI1, H1, AI2, H2, AI3]
-    result = ip.run_line_magic("ai", "reset")
+    ip.run_line_magic("ai", "reset")
     assert ai_magics.transcript == []

packages/jupyter-ai-magics/jupyter_ai_magics/utils.py

@@ -1,18 +1,18 @@
 import logging
-from typing import Dict, List, Literal, Optional, Tuple, Type, Union
+from typing import Literal, Optional, Union
 
 from importlib_metadata import entry_points
 from jupyter_ai_magics.aliases import MODEL_ID_ALIASES
 from jupyter_ai_magics.embedding_providers import BaseEmbeddingsProvider
 from jupyter_ai_magics.providers import BaseProvider
 
 Logger = Union[logging.Logger, logging.LoggerAdapter]
-LmProvidersDict = Dict[str, BaseProvider]
-EmProvidersDict = Dict[str, BaseEmbeddingsProvider]
+LmProvidersDict = dict[str, BaseProvider]
+EmProvidersDict = dict[str, BaseEmbeddingsProvider]
 AnyProvider = Union[BaseProvider, BaseEmbeddingsProvider]
-ProviderDict = Dict[str, AnyProvider]
-ProviderRestrictions = Dict[
-    Literal["allowed_providers", "blocked_providers"], Optional[List[str]]
+ProviderDict = dict[str, AnyProvider]
+ProviderRestrictions = dict[
+    Literal["allowed_providers", "blocked_providers"], Optional[list[str]]
 ]
 
 
@@ -80,8 +80,8 @@ def get_em_providers(
 
 
 def decompose_model_id(
-    model_id: str, providers: Dict[str, BaseProvider]
-) -> Tuple[str, str]:
+    model_id: str, providers: dict[str, BaseProvider]
+) -> tuple[str, str]:
     """Breaks down a model ID into a two-tuple (provider_id, local_model_id). Returns (None, None) if indeterminate."""
     if model_id in MODEL_ID_ALIASES:
         model_id = MODEL_ID_ALIASES[model_id]
@@ -104,15 +104,15 @@ def decompose_model_id(
 
 def get_lm_provider(
     model_id: str, lm_providers: LmProvidersDict
-) -> Tuple[str, Type[BaseProvider]]:
+) -> tuple[str, type[BaseProvider]]:
     """Gets a two-tuple (<local-model-id>, <provider-class>) specified by a
     global model ID."""
     return _get_provider(model_id, lm_providers)
 
 
 def get_em_provider(
     model_id: str, em_providers: EmProvidersDict
-) -> Tuple[str, Type[BaseEmbeddingsProvider]]:
+) -> tuple[str, type[BaseEmbeddingsProvider]]:
     """Gets a two-tuple (<local-model-id>, <provider-class>) specified by a
     global model ID."""
     return _get_provider(model_id, em_providers)