Add type hint for API classes

Author: chezou

tdclient/bulk_import_api.py

@@ -1,14 +1,25 @@
 #!/usr/bin/env python
 
+from __future__ import annotations
+
 import collections
 import contextlib
 import gzip
 import io
 import os
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
 
 import msgpack
 
-from .util import create_url
+if TYPE_CHECKING:
+    from contextlib import AbstractContextManager
+    from typing import IO
+
+    import urllib3
+
+from tdclient.types import BulkImportParams, BytesOrStream, DataFormat, FileLike
+from tdclient.util import create_url
 
 
 class BulkImportAPI:
@@ -17,7 +28,27 @@ class BulkImportAPI:
     This class is inherited by :class:`tdclient.api.API`.
     """
 
-    def create_bulk_import(self, name, db, table, params=None):
+    # Methods from API class
+    def get(
+        self, url: str, params: dict[str, Any] | None = None
+    ) -> AbstractContextManager[urllib3.BaseHTTPResponse]: ...
+    def post(
+        self, url: str, params: dict[str, Any] | None = None
+    ) -> AbstractContextManager[urllib3.BaseHTTPResponse]: ...
+    def put(
+        self, url: str, stream: BytesOrStream, size: int
+    ) -> AbstractContextManager[urllib3.BaseHTTPResponse]: ...
+    def raise_error(
+        self, msg: str, res: urllib3.BaseHTTPResponse, body: bytes | str
+    ) -> None: ...
+    def checked_json(self, body: bytes, required: list[str]) -> dict[str, Any]: ...
+    def _prepare_file(
+        self, file: FileLike, format: str, **kwargs: Any
+    ) -> IO[bytes]: ...
+
+    def create_bulk_import(
+        self, name: str, db: str, table: str, params: BulkImportParams | None = None
+    ) -> bool:
         """Enable bulk importing of data to the targeted database and table and stores
         it in the default resource pool. Default expiration for bulk import is 30days.
 
@@ -30,22 +61,24 @@ def create_bulk_import(self, name, db, table, params=None):
         Returns:
              True if succeeded
         """
-        params = {} if params is None else params
+        post_params = {} if params is None else dict(params)
         with self.post(
             create_url(
                 "/v3/bulk_import/create/{name}/{db}/{table}",
                 name=name,
                 db=db,
                 table=table,
             ),
-            params,
+            post_params,
         ) as res:
             code, body = res.status, res.read()
             if code != 200:
                 self.raise_error("Create bulk import failed", res, body)
             return True
 
-    def delete_bulk_import(self, name, params=None):
+    def delete_bulk_import(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> bool:
         """Delete the imported information with the specified name
 
         Args:
@@ -63,7 +96,7 @@ def delete_bulk_import(self, name, params=None):
                 self.raise_error("Delete bulk import failed", res, body)
             return True
 
-    def show_bulk_import(self, name):
+    def show_bulk_import(self, name: str) -> dict[str, Any]:
         """Show the details of the bulk import with the specified name
 
         Args:
@@ -78,7 +111,9 @@ def show_bulk_import(self, name):
             js = self.checked_json(body, ["status"])
             return js
 
-    def list_bulk_imports(self, params=None):
+    def list_bulk_imports(
+        self, params: dict[str, Any] | None = None
+    ) -> list[dict[str, Any]]:
         """Return the list of available bulk imports
         Args:
             params (dict, optional): Extra parameters.
@@ -93,7 +128,9 @@ def list_bulk_imports(self, params=None):
             js = self.checked_json(body, ["bulk_imports"])
             return js["bulk_imports"]
 
-    def list_bulk_import_parts(self, name, params=None):
+    def list_bulk_import_parts(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> list[str]:
         """Return the list of available parts uploaded through
         :func:`~BulkImportAPI.bulk_import_upload_part`.
 
@@ -114,7 +151,7 @@ def list_bulk_import_parts(self, name, params=None):
             return js["parts"]
 
     @staticmethod
-    def validate_part_name(part_name):
+    def validate_part_name(part_name: str) -> None:
         """Make sure the part_name is valid
 
         Args:
@@ -133,7 +170,9 @@ def validate_part_name(part_name):
         if 0 < part_name.find("/"):
             raise ValueError("part name must not contain '/': %s" % (repr(part_name)))
 
-    def bulk_import_upload_part(self, name, part_name, stream, size):
+    def bulk_import_upload_part(
+        self, name: str, part_name: str, stream: BytesOrStream, size: int
+    ) -> None:
         """Upload bulk import having the specified name and part in the path.
 
         Args:
@@ -156,7 +195,14 @@ def bulk_import_upload_part(self, name, part_name, stream, size):
             if code / 100 != 2:
                 self.raise_error("Upload a part failed", res, body)
 
-    def bulk_import_upload_file(self, name, part_name, format, file, **kwargs):
+    def bulk_import_upload_file(
+        self,
+        name: str,
+        part_name: str,
+        format: DataFormat,
+        file: FileLike,
+        **kwargs: Any,
+    ) -> None:
         """Upload a file with bulk import having the specified name.
 
         Args:
@@ -193,7 +239,9 @@ def bulk_import_upload_file(self, name, part_name, format, file, **kwargs):
             size = os.fstat(fp.fileno()).st_size
             return self.bulk_import_upload_part(name, part_name, fp, size)
 
-    def bulk_import_delete_part(self, name, part_name, params=None):
+    def bulk_import_delete_part(
+        self, name: str, part_name: str, params: dict[str, Any] | None = None
+    ) -> bool:
         """Delete the imported information with the specified name.
 
         Args:
@@ -218,7 +266,9 @@ def bulk_import_delete_part(self, name, part_name, params=None):
                 self.raise_error("Delete a part failed", res, body)
             return True
 
-    def freeze_bulk_import(self, name, params=None):
+    def freeze_bulk_import(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> bool:
         """Freeze the bulk import with the specified name.
 
         Args:
@@ -236,7 +286,9 @@ def freeze_bulk_import(self, name, params=None):
                 self.raise_error("Freeze bulk import failed", res, body)
             return True
 
-    def unfreeze_bulk_import(self, name, params=None):
+    def unfreeze_bulk_import(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> bool:
         """Unfreeze bulk_import with the specified name.
 
         Args:
@@ -254,7 +306,9 @@ def unfreeze_bulk_import(self, name, params=None):
                 self.raise_error("Unfreeze bulk import failed", res, body)
             return True
 
-    def perform_bulk_import(self, name, params=None):
+    def perform_bulk_import(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> str:
         """Execute a job to perform bulk import with the indicated priority using the
         resource pool if indicated, else it will use the account's default.
 
@@ -274,7 +328,9 @@ def perform_bulk_import(self, name, params=None):
             js = self.checked_json(body, ["job_id"])
             return str(js["job_id"])
 
-    def commit_bulk_import(self, name, params=None):
+    def commit_bulk_import(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> bool:
         """Commit the bulk import information having the specified name.
 
         Args:
@@ -292,7 +348,9 @@ def commit_bulk_import(self, name, params=None):
                 self.raise_error("Commit bulk import failed", res, body)
             return True
 
-    def bulk_import_error_records(self, name, params=None):
+    def bulk_import_error_records(
+        self, name: str, params: dict[str, Any] | None = None
+    ) -> Iterator[dict[str, Any]]:
         """List the records that have errors under the specified bulk import name.
 
         Args:

tdclient/client.py

@@ -5,14 +5,17 @@
 import datetime
 import json
 from collections.abc import Iterator
-from typing import Any
+from typing import Any, cast, Literal
+
 
 from tdclient import api, models
 from tdclient.types import (
     BulkImportParams,
+    BytesOrStream,
     DataFormat,
     ExportParams,
     FileLike,
+    Priority,
     ResultFormat,
     ResultParams,
     ScheduleParams,
@@ -254,7 +257,7 @@ def query(
         db_name: str,
         q: str,
         result_url: str | None = None,
-        priority: int | str | None = None,
+        priority: Priority | None = None,
         retry_limit: int | None = None,
         type: str = "hive",
         **kwargs: Any,
@@ -279,9 +282,11 @@ def query(
         # for compatibility, assume type is hive unless specifically specified
         if type not in ["hive", "pig", "impala", "presto", "trino"]:
             raise ValueError("The specified query type is not supported: %s" % (type))
+        # Cast type to expected literal since we've validated it
+        query_type = cast(Literal["hive", "presto", "trino", "bulkload"], type)
         job_id = self.api.query(
             q,
-            type=type,
+            type=query_type,
             db=db_name,
             result_url=result_url,
             priority=priority,
@@ -295,7 +300,7 @@ def jobs(
         _from: int | None = None,
         to: int | None = None,
         status: str | None = None,
-        conditions: str | None = None,
+        conditions: dict[str, Any] | None = None,
     ) -> list[models.Job]:
         """List jobs
 
@@ -304,7 +309,7 @@ def jobs(
             to (int, optional): Gets the Job up to the nth index in the list.
                 By default, the first 20 jobs in the list are displayed
             status (str, optional): Filter by given status. {"queued", "running", "success", "error"}
-            conditions (str, optional): Condition for ``TIMESTAMPDIFF()`` to search for slow queries.
+            conditions (dict[str, Any], optional): Condition for ``TIMESTAMPDIFF()`` to search for slow queries.
                 Avoid using this parameter as it can be dangerous.
 
         Returns:
@@ -334,7 +339,7 @@ def job_status(self, job_id: str | int) -> str:
         Returns:
              a string represents the status of the job ("success", "error", "killed", "queued", "running")
         """
-        return self.api.job_status(job_id)
+        return self.api.job_status(str(job_id))
 
     def job_result(self, job_id: str | int) -> list[Any]:
         """
@@ -344,7 +349,7 @@ def job_result(self, job_id: str | int) -> list[Any]:
         Returns:
              a list of each rows in result set
         """
-        return self.api.job_result(job_id)
+        return self.api.job_result(str(job_id))
 
     def job_result_each(self, job_id: str | int) -> Iterator[dict[str, Any]]:
         """
@@ -354,7 +359,7 @@ def job_result_each(self, job_id: str | int) -> Iterator[dict[str, Any]]:
         Returns:
              an iterator of result set
         """
-        for row in self.api.job_result_each(job_id):
+        for row in self.api.job_result_each(str(job_id)):
             yield row
 
     def job_result_format(
@@ -368,7 +373,7 @@ def job_result_format(
         Returns:
              a list of each rows in result set
         """
-        return self.api.job_result_format(job_id, format, header=header)
+        return self.api.job_result_format(str(job_id), format, header=header)
 
     def job_result_format_each(
         self,
@@ -393,7 +398,7 @@ def job_result_format_each(
              an iterator of rows in result set
         """
         for row in self.api.job_result_format_each(
-            job_id,
+            str(job_id),
             format,
             header=header,
             store_tmpfile=store_tmpfile,
@@ -414,17 +419,17 @@ def download_job_result(
         Returns:
              `True` if success
         """
-        return self.api.download_job_result(job_id, path, num_threads=num_threads)
+        return self.api.download_job_result(str(job_id), path, num_threads=num_threads)
 
-    def kill(self, job_id: str | int) -> str:
+    def kill(self, job_id: str | int) -> str | None:
         """
         Args:
             job_id (str): job id
 
         Returns:
              a string represents the status of killed job ("queued", "running")
         """
-        return self.api.kill(job_id)
+        return self.api.kill(str(job_id))
 
     def export_data(
         self,
@@ -582,7 +587,7 @@ def bulk_imports(self) -> list[models.BulkImport]:
         ]
 
     def bulk_import_upload_part(
-        self, name: str, part_name: str, bytes_or_stream: FileLike, size: int
+        self, name: str, part_name: str, bytes_or_stream: BytesOrStream, size: int
     ) -> None:
         """Upload a part to a bulk import session
 
@@ -849,7 +854,7 @@ def import_data(
         db_name: str,
         table_name: str,
         format: DataFormat,
-        bytes_or_stream: FileLike,
+        bytes_or_stream: BytesOrStream,
         size: int,
         unique_id: str | None = None,
     ) -> float: