Add type hints to 'eland.operations' and 'eland.ndframe'

eland/arithmetics.py

@@ -19,9 +19,11 @@
 from io import StringIO
 from typing import TYPE_CHECKING, Any, List, Union
 
-import numpy as np  # type: ignore
+import numpy as np
 
 if TYPE_CHECKING:
+    from numpy.typing import DTypeLike
+
     from .query_compiler import QueryCompiler
 
 
@@ -32,7 +34,7 @@ def value(self) -> str:
         pass
 
     @abstractmethod
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> "DTypeLike":
         pass
 
     @abstractmethod
@@ -52,7 +54,7 @@ def resolve(self) -> str:
         return self.value
 
     @property
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> "DTypeLike":
         return np.dtype(object)
 
     @property
@@ -64,7 +66,7 @@ def __repr__(self) -> str:
 
 
 class ArithmeticNumber(ArithmeticObject):
-    def __init__(self, value: Union[int, float], dtype: np.dtype):
+    def __init__(self, value: Union[int, float], dtype: "DTypeLike"):
         self._value = value
         self._dtype = dtype
 
@@ -76,7 +78,7 @@ def value(self) -> str:
         return f"{self._value}"
 
     @property
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> "DTypeLike":
         return self._dtype
 
     def __repr__(self) -> str:
@@ -89,8 +91,8 @@ class ArithmeticSeries(ArithmeticObject):
     """
 
     def __init__(
-        self, query_compiler: "QueryCompiler", display_name: str, dtype: np.dtype
-    ):
+        self, query_compiler: "QueryCompiler", display_name: str, dtype: "DTypeLike"
+    ) -> None:
         # type defs
         self._value: str
         self._tasks: List["ArithmeticTask"]
@@ -121,7 +123,7 @@ def value(self) -> str:
         return self._value
 
     @property
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> "DTypeLike":
         return self._dtype
 
     def __repr__(self) -> str:

eland/common.py

@@ -18,12 +18,24 @@
 import re
 import warnings
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
-import numpy as np  # type: ignore
 import pandas as pd  # type: ignore
 from elasticsearch import Elasticsearch
 
+if TYPE_CHECKING:
+    from numpy.typing import DTypeLike
+
 # Default number of rows displayed (different to pandas where ALL could be displayed)
 DEFAULT_NUM_ROWS_DISPLAYED = 60
 DEFAULT_CHUNK_SIZE = 10000
@@ -42,7 +54,7 @@
 
 
 def build_pd_series(
-    data: Dict[str, Any], dtype: Optional[np.dtype] = None, **kwargs: Any
+    data: Dict[str, Any], dtype: Optional["DTypeLike"] = None, **kwargs: Any
 ) -> pd.Series:
     """Builds a pd.Series while squelching the warning
     for unspecified dtype on empty series
@@ -88,7 +100,7 @@ def from_string(order: str) -> "SortOrder":
 
 
 def elasticsearch_date_to_pandas_date(
-    value: Union[int, str], date_format: Optional[str]
+    value: Union[int, str, float], date_format: Optional[str]
 ) -> pd.Timestamp:
     """
     Given a specific Elasticsearch format for a date datatype, returns the
@@ -98,7 +110,7 @@ def elasticsearch_date_to_pandas_date(
 
     Parameters
     ----------
-    value: Union[int, str]
+    value: Union[int, str, float]
         The date value.
     date_format: str
         The Elasticsearch date format (ex. 'epoch_millis', 'epoch_second', etc.)

eland/conftest.py

@@ -15,9 +15,11 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+from typing import Any, Dict
+
 import numpy as np
-import pandas as pd
-import pytest
+import pandas as pd  # type: ignore
+import pytest  # type: ignore
 
 import eland as ed
 
@@ -28,7 +30,7 @@
 
 
 @pytest.fixture(autouse=True)
-def add_imports(doctest_namespace):
+def add_imports(doctest_namespace: Dict[str, Any]) -> None:
     doctest_namespace["np"] = np
     doctest_namespace["pd"] = pd
     doctest_namespace["ed"] = ed