Modernize type annotations in atomicity modules.

PiperOrigin-RevId: 787073942

Author: dicentra13

checkpoint/orbax/checkpoint/_src/path/atomicity.py

@@ -56,7 +56,7 @@
 import pickle
 import threading
 import time
-from typing import Awaitable, Optional, Protocol, Sequence, Type, TypeVar
+from typing import Awaitable, Protocol, Sequence, TypeVar
 
 from absl import logging
 from etils import epath
@@ -101,7 +101,7 @@ def __call__(
       path: epath.Path,
       parents: bool = False,
       exist_ok: bool = False,
-      mode: Optional[int] = None,
+      mode: int | None = None,
       **kwargs,
   ) -> Awaitable[None]:
     """Creates the directory at path."""
@@ -112,10 +112,10 @@ async def _create_tmp_directory(
     async_makedir_func: AsyncMakeDirFunc,
     tmp_dir: epath.Path,
     *,
-    path_permission_mode: Optional[int] = None,
-    checkpoint_metadata_store: Optional[
-        checkpoint_metadata.MetadataStore
-    ] = None,
+    path_permission_mode: int | None = None,
+    checkpoint_metadata_store: (
+        checkpoint_metadata.MetadataStore | None
+    ) = None,
     **kwargs,
 ) -> epath.Path:
   """Creates a non-deterministic tmp directory for saving for given `final_dir`.
@@ -181,7 +181,7 @@ def _get_tmp_directory(final_path: epath.Path) -> epath.Path:
   )
 
 
-def _get_tmp_directory_pattern(final_path_name: Optional[str] = None) -> str:
+def _get_tmp_directory_pattern(final_path_name: str | None = None) -> str:
   suffix = r'\.orbax-checkpoint-tmp'
   if final_path_name is None:
     return '(.+)' + suffix
@@ -197,10 +197,10 @@ def __init__(
       temporary_path: epath.Path,
       final_path: epath.Path,
       *,
-      checkpoint_metadata_store: Optional[
-          checkpoint_metadata.MetadataStore
-      ] = None,
-      file_options: Optional[options_lib.FileOptions] = None,
+      checkpoint_metadata_store: (
+          checkpoint_metadata.MetadataStore | None
+      ) = None,
+      file_options: options_lib.FileOptions | None = None,
   ):
     self._tmp_path = temporary_path
     self._final_path = final_path
@@ -269,7 +269,7 @@ def to_bytes(self) -> bytes:
 
   @classmethod
   def from_bytes(
-      cls: Type['ReadOnlyTemporaryPath'],
+      cls: type['ReadOnlyTemporaryPath'],
       data: bytes,
   ) -> ReadOnlyTemporaryPath:
     """Deserializes the object from bytes.
@@ -291,10 +291,10 @@ def from_final(
       cls,
       final_path: epath.Path,
       *,
-      checkpoint_metadata_store: Optional[
-          checkpoint_metadata.MetadataStore
-      ] = None,
-      file_options: Optional[options_lib.FileOptions] = None,
+      checkpoint_metadata_store: (
+          checkpoint_metadata.MetadataStore | None
+      ) = None,
+      file_options: options_lib.FileOptions | None = None,
   ) -> ReadOnlyTemporaryPath:
     """Not implemented for ReadOnlyTemporaryPath."""
     raise NotImplementedError(
@@ -324,10 +324,10 @@ def from_final(
       cls,
       final_path: epath.Path,
       *,
-      checkpoint_metadata_store: Optional[
-          checkpoint_metadata.MetadataStore
-      ] = None,
-      file_options: Optional[options_lib.FileOptions] = None,
+      checkpoint_metadata_store: (
+          checkpoint_metadata.MetadataStore | None
+      ) = None,
+      file_options: options_lib.FileOptions | None = None,
   ) -> AtomicRenameTemporaryPath:
     return cls(
         _get_tmp_directory(final_path),
@@ -399,10 +399,10 @@ def from_final(
       cls,
       final_path: epath.Path,
       *,
-      checkpoint_metadata_store: Optional[
-          checkpoint_metadata.MetadataStore
-      ] = None,
-      file_options: Optional[options_lib.FileOptions] = None,
+      checkpoint_metadata_store: (
+          checkpoint_metadata.MetadataStore | None
+      ) = None,
+      file_options: options_lib.FileOptions | None = None,
   ) -> CommitFileTemporaryPath:
     return cls(
         final_path,
@@ -465,9 +465,7 @@ def finalize(
 async def create_all(
     paths: Sequence[atomicity_types.TemporaryPath],
     *,
-    multiprocessing_options: Optional[
-        options_lib.MultiprocessingOptions
-    ] = None,
+    multiprocessing_options: options_lib.MultiprocessingOptions | None = None,
 ):
   """Creates all temporary paths in parallel."""
   start = time.time()
@@ -511,9 +509,7 @@ def create_all_async(
     paths: Sequence[atomicity_types.TemporaryPath],
     completion_signals: Sequence[synchronization.HandlerAwaitableSignal],
     *,
-    multiprocessing_options: Optional[
-        options_lib.MultiprocessingOptions
-    ] = None,
+    multiprocessing_options: options_lib.MultiprocessingOptions | None = None,
     subdirectories: Sequence[str] | None = None,
 ) -> future.Future:
   """Creates all temporary paths in parallel asynchronously.

checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py

@@ -19,7 +19,6 @@
 than we would want to introduce into the base `atomicity` module.
 """
 
-from typing import Type
 from etils import epath
 from orbax.checkpoint._src.path import atomicity
 from orbax.checkpoint._src.path import atomicity_types
@@ -28,7 +27,7 @@
 
 def get_item_default_temporary_path_class(
     final_path: epath.Path,
-) -> Type[atomicity_types.TemporaryPath]:
+) -> type[atomicity_types.TemporaryPath]:
   """Returns the default temporary path class for a given sub-item path."""
   if step_lib.is_gcs_path(final_path):
     return atomicity.CommitFileTemporaryPath
@@ -38,7 +37,7 @@ def get_item_default_temporary_path_class(
 
 def get_default_temporary_path_class(
     final_path: epath.Path,
-) -> Type[atomicity_types.TemporaryPath]:
+) -> type[atomicity_types.TemporaryPath]:
   """Returns the default temporary path class for a given checkpoint path."""
   if step_lib.is_gcs_path(final_path):
     return atomicity.CommitFileTemporaryPath

checkpoint/orbax/checkpoint/_src/path/atomicity_types.py

@@ -21,7 +21,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Optional, Protocol
+from typing import Protocol
 
 from etils import epath
 from orbax.checkpoint import options as options_lib
@@ -49,10 +49,10 @@ def from_final(
       cls,
       final_path: epath.Path,
       *,
-      checkpoint_metadata_store: Optional[
-          checkpoint_metadata.MetadataStore
-      ] = None,
-      file_options: Optional[options_lib.FileOptions] = None,
+      checkpoint_metadata_store: (
+          checkpoint_metadata.MetadataStore | None
+      ) = None,
+      file_options: options_lib.FileOptions | None = None,
   ) -> TemporaryPath:
     """Creates a TemporaryPath from a final path."""
     ...