Implement PEP 604 with future-annotations

Author: dgarros

infrahub_sdk/_importer.py

@@ -3,7 +3,7 @@
 import importlib
 import sys
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from .exceptions import ModuleImportError
 
@@ -13,9 +13,7 @@
 module_mtime_cache: dict[str, float] = {}
 
 
-def import_module(
-    module_path: Path, import_root: Optional[str] = None, relative_path: Optional[str] = None
-) -> ModuleType:
+def import_module(module_path: Path, import_root: str | None = None, relative_path: str | None = None) -> ModuleType:
     """Imports a python module.
 
     Args:

infrahub_sdk/analyzer.py

@@ -1,4 +1,6 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import Any
 
 from graphql import (
     DocumentNode,
@@ -19,23 +21,23 @@ class GraphQLQueryVariable(BaseModel):
     name: str
     type: str
     required: bool = False
-    default_value: Optional[Any] = None
+    default_value: Any | None = None
 
 
 class GraphQLOperation(BaseModel):
-    name: Optional[str] = None
+    name: str | None = None
     operation_type: OperationType
 
 
 class GraphQLQueryAnalyzer:
-    def __init__(self, query: str, schema: Optional[GraphQLSchema] = None):
+    def __init__(self, query: str, schema: GraphQLSchema | None = None):
         self.query: str = query
-        self.schema: Optional[GraphQLSchema] = schema
+        self.schema: GraphQLSchema | None = schema
         self.document: DocumentNode = parse(self.query)
-        self._fields: Optional[dict] = None
+        self._fields: dict | None = None
 
     @property
-    def is_valid(self) -> tuple[bool, Optional[list[GraphQLError]]]:
+    def is_valid(self) -> tuple[bool, list[GraphQLError] | None]:
         if self.schema is None:
             return False, [GraphQLError("Schema is not provided")]
 

infrahub_sdk/batch.py

@@ -1,28 +1,31 @@
+from __future__ import annotations
+
 import asyncio
-from collections.abc import AsyncGenerator, Awaitable
+from collections.abc import AsyncGenerator, Awaitable, Generator
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-from typing import Any, Callable, Generator, Optional
+from typing import TYPE_CHECKING, Any, Callable
 
-from .node import InfrahubNode, InfrahubNodeSync
+if TYPE_CHECKING:
+    from .node import InfrahubNode, InfrahubNodeSync
 
 
 @dataclass
 class BatchTask:
     task: Callable[[Any], Awaitable[Any]]
     args: tuple[Any, ...]
     kwargs: dict[str, Any]
-    node: Optional[Any] = None
+    node: Any | None = None
 
 
 @dataclass
 class BatchTaskSync:
     task: Callable[..., Any]
     args: tuple[Any, ...]
     kwargs: dict[str, Any]
-    node: Optional[InfrahubNodeSync] = None
+    node: InfrahubNodeSync | None = None
 
-    def execute(self, return_exceptions: bool = False) -> tuple[Optional[InfrahubNodeSync], Any]:
+    def execute(self, return_exceptions: bool = False) -> tuple[InfrahubNodeSync | None, Any]:
         """Executes the stored task."""
         result = None
         try:
@@ -37,7 +40,7 @@ def execute(self, return_exceptions: bool = False) -> tuple[Optional[InfrahubNod
 
 async def execute_batch_task_in_pool(
     task: BatchTask, semaphore: asyncio.Semaphore, return_exceptions: bool = False
-) -> tuple[Optional[InfrahubNode], Any]:
+) -> tuple[InfrahubNode | None, Any]:
     async with semaphore:
         try:
             result = await task.task(*task.args, **task.kwargs)
@@ -52,7 +55,7 @@ async def execute_batch_task_in_pool(
 class InfrahubBatch:
     def __init__(
         self,
-        semaphore: Optional[asyncio.Semaphore] = None,
+        semaphore: asyncio.Semaphore | None = None,
         max_concurrent_execution: int = 5,
         return_exceptions: bool = False,
     ):
@@ -64,7 +67,7 @@ def __init__(
     def num_tasks(self) -> int:
         return len(self._tasks)
 
-    def add(self, *args: Any, task: Callable, node: Optional[Any] = None, **kwargs: Any) -> None:
+    def add(self, *args: Any, task: Callable, node: Any | None = None, **kwargs: Any) -> None:
         self._tasks.append(BatchTask(task=task, node=node, args=args, kwargs=kwargs))
 
     async def execute(self) -> AsyncGenerator:
@@ -96,10 +99,10 @@ def __init__(self, max_concurrent_execution: int = 5, return_exceptions: bool =
     def num_tasks(self) -> int:
         return len(self._tasks)
 
-    def add(self, *args: Any, task: Callable[..., Any], node: Optional[Any] = None, **kwargs: Any) -> None:
+    def add(self, *args: Any, task: Callable[..., Any], node: Any | None = None, **kwargs: Any) -> None:
         self._tasks.append(BatchTaskSync(task=task, node=node, args=args, kwargs=kwargs))
 
-    def execute(self) -> Generator[tuple[Optional[InfrahubNodeSync], Any], None, None]:
+    def execute(self) -> Generator[tuple[InfrahubNodeSync | None, Any], None, None]:
         with ThreadPoolExecutor(max_workers=self.max_concurrent_execution) as executor:
             futures = [executor.submit(task.execute, return_exceptions=self.return_exceptions) for task in self._tasks]
             for future in futures:

infrahub_sdk/branch.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, Literal, overload
 from urllib.parse import urlencode
 
 from pydantic import BaseModel
@@ -17,11 +17,11 @@
 class BranchData(BaseModel):
     id: str
     name: str
-    description: Optional[str] = None
+    description: str | None = None
     sync_with_git: bool
     is_default: bool
     has_schema_changes: bool
-    origin_branch: Optional[str] = None
+    origin_branch: str | None = None
     branched_from: str
 
 
@@ -50,11 +50,11 @@ class InfraHubBranchManagerBase:
     @classmethod
     def generate_diff_data_url(
         cls,
-        client: Union[InfrahubClient, InfrahubClientSync],
+        client: InfrahubClient | InfrahubClientSync,
         branch_name: str,
         branch_only: bool = True,
-        time_from: Optional[str] = None,
-        time_to: Optional[str] = None,
+        time_from: str | None = None,
+        time_to: str | None = None,
     ) -> str:
         """Generate the URL for the diff_data function."""
         url = f"{client.address}/api/diff/data"
@@ -80,7 +80,7 @@ async def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: Literal[True] = True,
-        background_execution: Optional[bool] = False,
+        background_execution: bool | None = False,
     ) -> BranchData: ...
 
     @overload
@@ -90,7 +90,7 @@ async def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: Literal[False] = False,
-        background_execution: Optional[bool] = False,
+        background_execution: bool | None = False,
     ) -> str: ...
 
     async def create(
@@ -99,8 +99,8 @@ async def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: bool = True,
-        background_execution: Optional[bool] = False,
-    ) -> Union[BranchData, str]:
+        background_execution: bool | None = False,
+    ) -> BranchData | str:
         if background_execution is not None:
             warnings.warn(
                 "`background_execution` is deprecated, please use `wait_until_completion` instead.",
@@ -206,8 +206,8 @@ async def diff_data(
         self,
         branch_name: str,
         branch_only: bool = True,
-        time_from: Optional[str] = None,
-        time_to: Optional[str] = None,
+        time_from: str | None = None,
+        time_to: str | None = None,
     ) -> dict[Any, Any]:
         url = self.generate_diff_data_url(
             client=self.client,
@@ -251,7 +251,7 @@ def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: Literal[True] = True,
-        background_execution: Optional[bool] = False,
+        background_execution: bool | None = False,
     ) -> BranchData: ...
 
     @overload
@@ -261,7 +261,7 @@ def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: Literal[False] = False,
-        background_execution: Optional[bool] = False,
+        background_execution: bool | None = False,
     ) -> str: ...
 
     def create(
@@ -270,8 +270,8 @@ def create(
         sync_with_git: bool = True,
         description: str = "",
         wait_until_completion: bool = True,
-        background_execution: Optional[bool] = False,
-    ) -> Union[BranchData, str]:
+        background_execution: bool | None = False,
+    ) -> BranchData | str:
         if background_execution is not None:
             warnings.warn(
                 "`background_execution` is deprecated, please use `wait_until_completion` instead.",
@@ -313,8 +313,8 @@ def diff_data(
         self,
         branch_name: str,
         branch_only: bool = True,
-        time_from: Optional[str] = None,
-        time_to: Optional[str] = None,
+        time_from: str | None = None,
+        time_to: str | None = None,
     ) -> dict[Any, Any]:
         url = self.generate_diff_data_url(
             client=self.client,

infrahub_sdk/checks.py

@@ -5,7 +5,7 @@
 import os
 import warnings
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import ujson
 from git.repo import Repo
@@ -33,20 +33,20 @@ class InfrahubCheckInitializer(BaseModel):
 
 
 class InfrahubCheck:
-    name: Optional[str] = None
+    name: str | None = None
     query: str = ""
     timeout: int = 10
 
     def __init__(
         self,
-        branch: Optional[str] = None,
+        branch: str | None = None,
         root_directory: str = "",
-        output: Optional[str] = None,
-        initializer: Optional[InfrahubCheckInitializer] = None,
-        params: Optional[dict] = None,
-        client: Optional[InfrahubClient] = None,
+        output: str | None = None,
+        initializer: InfrahubCheckInitializer | None = None,
+        params: dict | None = None,
+        client: InfrahubClient | None = None,
     ):
-        self.git: Optional[Repo] = None
+        self.git: Repo | None = None
         self.initializer = initializer or InfrahubCheckInitializer()
 
         self.logs: list[dict[str, Any]] = []
@@ -82,7 +82,7 @@ def client(self, value: InfrahubClient) -> None:
         self._client = value
 
     @classmethod
-    async def init(cls, client: Optional[InfrahubClient] = None, *args: Any, **kwargs: Any) -> InfrahubCheck:
+    async def init(cls, client: InfrahubClient | None = None, *args: Any, **kwargs: Any) -> InfrahubCheck:
         """Async init method, If an existing InfrahubClient client hasn't been provided, one will be created automatically."""
         warnings.warn(
             "InfrahubCheck.init has been deprecated and will be removed in the version in Infrahub SDK 2.0.0",
@@ -101,7 +101,7 @@ def errors(self) -> list[dict[str, Any]]:
         return [log for log in self.logs if log["level"] == "ERROR"]
 
     def _write_log_entry(
-        self, message: str, level: str, object_id: Optional[str] = None, object_type: Optional[str] = None
+        self, message: str, level: str, object_id: str | None = None, object_type: str | None = None
     ) -> None:
         log_message = {"level": level, "message": message, "branch": self.branch_name}
         if object_id:
@@ -113,10 +113,10 @@ def _write_log_entry(
         if self.output == "stdout":
             print(ujson.dumps(log_message))
 
-    def log_error(self, message: str, object_id: Optional[str] = None, object_type: Optional[str] = None) -> None:
+    def log_error(self, message: str, object_id: str | None = None, object_type: str | None = None) -> None:
         self._write_log_entry(message=message, level="ERROR", object_id=object_id, object_type=object_type)
 
-    def log_info(self, message: str, object_id: Optional[str] = None, object_type: Optional[str] = None) -> None:
+    def log_info(self, message: str, object_id: str | None = None, object_type: str | None = None) -> None:
         self._write_log_entry(message=message, level="INFO", object_id=object_id, object_type=object_type)
 
     @property
@@ -155,7 +155,7 @@ async def collect_data(self) -> dict:
 
         return await self.client.query_gql_query(name=self.query, branch_name=self.branch_name, variables=self.params)
 
-    async def run(self, data: Optional[dict] = None) -> bool:
+    async def run(self, data: dict | None = None) -> bool:
         """Execute the check after collecting the data from the GraphQL query.
         The result of the check is determined based on the presence or not of ERROR log messages."""
 
@@ -179,7 +179,7 @@ async def run(self, data: Optional[dict] = None) -> bool:
 
 
 def get_check_class_instance(
-    check_config: InfrahubCheckDefinitionConfig, search_path: Optional[Path] = None
+    check_config: InfrahubCheckDefinitionConfig, search_path: Path | None = None
 ) -> InfrahubCheck:
     if check_config.file_path.is_absolute() or search_path is None:
         search_location = check_config.file_path