Move typing-only imports under if TYPE_CHECKING

Author: aversey

python/hsfs/connection.py

@@ -17,9 +17,9 @@
 
 import importlib.util
 import os
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
-from hsfs import client, engine, feature_store, usage, util
+from hsfs import client, engine, usage, util
 from hsfs.core import (
     feature_store_api,
     hosts_api,
@@ -32,6 +32,10 @@
 from requests.exceptions import ConnectionError
 
 
+if TYPE_CHECKING:
+    from hsfs import feature_store
+
+
 AWS_DEFAULT_REGION = "default"
 HOPSWORKS_PORT_DEFAULT = 443
 SECRETS_STORE_DEFAULT = "parameterstore"

python/hsfs/constructor/fs_query.py

@@ -15,13 +15,17 @@
 #
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import humps
 from hsfs import engine
 from hsfs.constructor import external_feature_group_alias, hudi_feature_group_alias
 
 
+if TYPE_CHECKING:
+    import pyspark
+
+
 class FsQuery:
     def __init__(
         self,
@@ -96,7 +100,7 @@ def hudi_cached_feature_groups(
     def register_external(
         self,
         spine: Optional[
-            Union[TypeVar("pyspark.sql.DataFrame"), TypeVar("pyspark.RDD")]
+            Union[pyspark.sql.DataFrame, pyspark.RDD]
         ] = None,
     ) -> None:
         if self._on_demand_fg_aliases is None:

python/hsfs/constructor/query.py

@@ -18,20 +18,25 @@
 import json
 import warnings
 from datetime import date, datetime
-from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import humps
-import numpy as np
-import pandas as pd
-from hsfs import engine, storage_connector, util
+from hsfs import engine, util
 from hsfs import feature_group as fg_mod
 from hsfs.client.exceptions import FeatureStoreException
 from hsfs.constructor import join
 from hsfs.constructor.filter import Filter, Logic
-from hsfs.constructor.fs_query import FsQuery
 from hsfs.core import arrow_flight_client, query_constructor_api, storage_connector_api
 from hsfs.decorators import typechecked
-from hsfs.feature import Feature
+
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+    import pyspark
+    from hsfs import storage_connector
+    from hsfs.constructor.fs_query import FsQuery
+    from hsfs.feature import Feature
 
 
 @typechecked
@@ -59,7 +64,7 @@ def __init__(
             fg_mod.ExternalFeatureGroup,
             fg_mod.SpineGroup,
         ],
-        left_features: List[Union[str, "Feature", Dict]],
+        left_features: List[Union[str, Feature, Dict]],
         feature_store_name: Optional[str] = None,
         feature_store_id: Optional[int] = None,
         left_feature_group_start_time: Optional[Union[str, int, date, datetime]] = None,
@@ -80,14 +85,14 @@ def __init__(
         self._query_constructor_api: "query_constructor_api.QueryConstructorApi" = (
             query_constructor_api.QueryConstructorApi()
         )
-        self._storage_connector_api: "storage_connector_api.StorageConnectorApi" = (
+        self._storage_connector_api: storage_connector_api.StorageConnectorApi = (
             storage_connector_api.StorageConnectorApi()
         )
 
     def _prep_read(
         self, online: bool, read_options: Dict[str, Any]
     ) -> Tuple[
-        Union[str, Dict[str, Any]], Optional["storage_connector.StorageConnector"]
+        Union[str, Dict[str, Any]], Optional[storage_connector.StorageConnector]
     ]:
         self._check_read_supported(online)
         fs_query = self._query_constructor_api.construct_query(self)
@@ -141,8 +146,8 @@ def read(
         pd.DataFrame,
         np.ndarray,
         List[List[Any]],
-        TypeVar("pyspark.sql.DataFrame"),
-        TypeVar("pyspark.RDD"),
+        pyspark.sql.DataFrame,
+        pyspark.RDD,
     ]:
         """Read the specified query into a DataFrame.
 
@@ -595,7 +600,7 @@ def to_string(self, online: bool = False, arrow_flight: bool = False) -> str:
         return self._to_string(fs_query, online, arrow_flight)
 
     def _to_string(
-        self, fs_query: "FsQuery", online: bool = False, asof: bool = False
+        self, fs_query: FsQuery, online: bool = False, asof: bool = False
     ) -> str:
         if online:
             return fs_query.query_online
@@ -633,7 +638,7 @@ def left_feature_group_end_time(
     ) -> None:
         self._left_feature_group_end_time = left_feature_group_end_time
 
-    def append_feature(self, feature: Union[str, "Feature"]) -> "Query":
+    def append_feature(self, feature: Union[str, Feature]) -> "Query":
         """
         Append a feature to the query.
 
@@ -699,7 +704,7 @@ def _get_feature_by_name(
         self,
         feature_name: str,
     ) -> Tuple[
-        "Feature",
+        Feature,
         Optional[str],
         Union[
             fg_mod.FeatureGroup,
@@ -785,7 +790,7 @@ def filters(self) -> Optional[Logic]:
         return filters
 
     @property
-    def features(self) -> List["Feature"]:
+    def features(self) -> List[Feature]:
         """List of all features in the query"""
         features = []
         for feat in self._left_features:
@@ -797,7 +802,7 @@ def features(self) -> List["Feature"]:
 
         return features
 
-    def get_feature(self, feature_name: str) -> "Feature":
+    def get_feature(self, feature_name: str) -> Feature:
         """
         Get a feature by name.
 

python/hsfs/core/vector_db_client.py

@@ -17,14 +17,17 @@
 
 import base64
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import hsfs
 from hsfs.client.exceptions import FeatureStoreException, VectorDatabaseException
 from hsfs.constructor.filter import Filter, Logic
 from hsfs.constructor.join import Join
 from hsfs.core.opensearch import OpenSearchClientSingleton
-from hsfs.feature import Feature
+
+
+if TYPE_CHECKING:
+    from hsfs.feature import Feature
 
 
 class VectorDbClient:

python/hsfs/feature.py

@@ -17,15 +17,18 @@
 
 import json
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-import hsfs
 import humps
 from hsfs import util
 from hsfs.constructor import filter
 from hsfs.decorators import typechecked
 
 
+if TYPE_CHECKING:
+    import hsfs
+
+
 @typechecked
 class Feature:
     """Metadata object representing a feature in a feature group in the Feature Store.
@@ -49,9 +52,9 @@ def __init__(
         feature_group_id: Optional[int] = None,
         feature_group: Optional[
             Union[
-                "hsfs.feature_group.FeatureGroup",
-                "hsfs.feature_group.ExternalFeatureGroup",
-                "hsfs.feature_group.SpineGroup",
+                hsfs.feature_group.FeatureGroup,
+                hsfs.feature_group.ExternalFeatureGroup,
+                hsfs.feature_group.SpineGroup,
             ]
         ] = None,
         **kwargs,
@@ -216,41 +219,41 @@ def _get_filter_value(self, value: Any) -> Any:
         else:
             return value
 
-    def __lt__(self, other: Any) -> "filter.Filter":
+    def __lt__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.LT,
                              self._get_filter_value(other))
 
-    def __le__(self, other: Any) -> "filter.Filter":
+    def __le__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.LE,
                              self._get_filter_value(other))
 
-    def __eq__(self, other: Any) -> "filter.Filter":
+    def __eq__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.EQ,
                              self._get_filter_value(other))
 
-    def __ne__(self, other: Any) -> "filter.Filter":
+    def __ne__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.NE,
                              self._get_filter_value(other))
 
-    def __ge__(self, other: Any) -> "filter.Filter":
+    def __ge__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.GE,
                              self._get_filter_value(other))
 
-    def __gt__(self, other: Any) -> "filter.Filter":
+    def __gt__(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.GT,
                              self._get_filter_value(other))
 
-    def contains(self, other: Union[str, List[Any]]) -> "filter.Filter":
+    def contains(self, other: Union[str, List[Any]]) -> filter.Filter:
         """
         !!! warning "Deprecated"
             `contains` method is deprecated. Use `isin` instead.
         """
         return self.isin(other)
 
-    def isin(self, other: Union[str, List[Any]]) -> "filter.Filter":
+    def isin(self, other: Union[str, List[Any]]) -> filter.Filter:
         return filter.Filter(self, filter.Filter.IN, json.dumps(other))
 
-    def like(self, other: Any) -> "filter.Filter":
+    def like(self, other: Any) -> filter.Filter:
         return filter.Filter(self, filter.Filter.LK, other)
 
     def __str__(self) -> str: