--unsafe-fixes

Author: Borda

src/litmodels/integrations/mixins.py

@@ -4,7 +4,7 @@
 import warnings
 from abc import ABC
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from lightning_utilities.core.rank_zero import rank_zero_warn
 
@@ -43,7 +43,7 @@ def download_model(
 
     def _setup(
         self, name: Optional[str] = None, temp_folder: Union[str, Path, None] = None
-    ) -> Tuple[str, str, Union[str, Path]]:
+    ) -> tuple[str, str, Union[str, Path]]:
         """Parse and validate the model name and temporary folder."""
         if name is None:
             name = model_name = self.__class__.__name__
@@ -56,7 +56,7 @@ def _setup(
         return name, model_name, temp_folder
 
     def _upload_model_files(
-        self, name: str, path: Union[str, Path, List[Union[str, Path]]], metadata: Optional[dict] = None
+        self, name: str, path: Union[str, Path, list[Union[str, Path]]], metadata: Optional[dict] = None
     ) -> None:
         """Upload the model files to the registry."""
         if not metadata:

src/litmodels/io/cloud.py

@@ -3,7 +3,7 @@
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.models import _extend_model_name_with_teamspace, _parse_org_teamspace_model_version
@@ -45,11 +45,11 @@ def _print_model_link(name: str, verbose: Union[bool, int]) -> None:
 
 def upload_model_files(
     name: str,
-    path: Union[str, Path, List[Union[str, Path]]],
+    path: Union[str, Path, list[Union[str, Path]]],
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     verbose: Union[bool, int] = 1,
-    metadata: Optional[Dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
     """Upload a local checkpoint file to the model store.
 
@@ -83,7 +83,7 @@ def download_model_files(
     name: str,
     download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
-) -> Union[str, List[str]]:
+) -> Union[str, list[str]]:
     """Download a checkpoint from the model store.
 
     Args:
@@ -103,7 +103,7 @@ def download_model_files(
     )
 
 
-def _list_available_teamspaces() -> Dict[str, dict]:
+def _list_available_teamspaces() -> dict[str, dict]:
     """List available teamspaces for the authenticated user.
 
     Returns:

src/litmodels/io/gateway.py

@@ -1,7 +1,7 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from litmodels.io.cloud import download_model_files, upload_model_files
 from litmodels.io.utils import _KERAS_AVAILABLE, _PYTORCH_AVAILABLE, dump_pickle, load_pickle
@@ -22,7 +22,7 @@ def upload_model(
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     verbose: Union[bool, int] = 1,
-    metadata: Optional[Dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
     """Upload a checkpoint to the model store.
 
@@ -60,7 +60,7 @@ def save_model(
     cloud_account: Optional[str] = None,
     staging_dir: Optional[str] = None,
     verbose: Union[bool, int] = 1,
-    metadata: Optional[Dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
 ) -> "UploadedModelInfo":
     """Upload a checkpoint to the model store.
 
@@ -119,7 +119,7 @@ def download_model(
     name: str,
     download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
-) -> Union[str, List[str]]:
+) -> Union[str, list[str]]:
     """Download a checkpoint from the model store.
 
     Args: