More type hint with strict mode

.github/workflows/codeql-analysis.yml

@@ -1,3 +1,5 @@
+# Comment out due to conflict with the default CodeQL workflow
+#
 # For most projects, this workflow file will not need changing; you simply need
 # to commit it to your repository.
 #
@@ -9,64 +11,64 @@
 # the `language` matrix defined below to confirm you have the correct set of
 # supported CodeQL languages.
 #
-name: "CodeQL"
+# name: "CodeQL"
 
-on:
-  push:
-    branches: [ master ]
-  pull_request:
-    # The branches below must be a subset of the branches above
-    branches: [ master ]
-  schedule:
-    - cron: '38 8 * * 3'
+# on:
+#   push:
+#     branches: [ master ]
+#   pull_request:
+#     # The branches below must be a subset of the branches above
+#     branches: [ master ]
+#   schedule:
+#     - cron: '38 8 * * 3'
 
-jobs:
-  analyze:
-    name: Analyze
-    runs-on: ubuntu-latest
-    permissions:
-      actions: read
-      contents: read
-      security-events: write
+# jobs:
+#   analyze:
+#     name: Analyze
+#     runs-on: ubuntu-latest
+#     permissions:
+#       actions: read
+#       contents: read
+#       security-events: write
 
-    strategy:
-      fail-fast: false
-      matrix:
-        language: [ 'python' ]
-        # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ]
-        # Learn more about CodeQL language support at https://aka.ms/codeql-docs/language-support
+#     strategy:
+#       fail-fast: false
+#       matrix:
+#         language: [ 'python' ]
+#         # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ]
+#         # Learn more about CodeQL language support at https://aka.ms/codeql-docs/language-support
 
-    steps:
-    - name: Checkout repository
-      uses: actions/checkout@v5
+#     steps:
+#     - name: Checkout repository
+#       uses: actions/checkout@v5
 
-    # Initializes the CodeQL tools for scanning.
-    - name: Initialize CodeQL
-      uses: github/codeql-action/init@v2
-      with:
-        languages: ${{ matrix.language }}
-        # If you wish to specify custom queries, you can do so here or in a config file.
-        # By default, queries listed here will override any specified in a config file.
-        # Prefix the list here with "+" to use these queries and those in the config file.
+#     # Initializes the CodeQL tools for scanning.
+#     - name: Initialize CodeQL
+#       uses: github/codeql-action/init@v2
+#       with:
+#         languages: ${{ matrix.language }}
+#         # If you wish to specify custom queries, you can do so here or in a config file.
+#         # By default, queries listed here will override any specified in a config file.
+#         # Prefix the list here with "+" to use these queries and those in the config file.
 
-        # Details on CodeQL's query packs refer to : https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
-        # queries: security-extended,security-and-quality
+#         # Details on CodeQL's query packs refer to : https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
+#         # queries: security-extended,security-and-quality
 
 
-    # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
-    # If this step fails, then you should remove it and run the build manually (see below)
-    - name: Autobuild
-      uses: github/codeql-action/autobuild@v2
+#     # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
+#     # If this step fails, then you should remove it and run the build manually (see below)
+#     - name: Autobuild
+#       uses: github/codeql-action/autobuild@v2
 
-    # ℹ️ Command-line programs to run using the OS shell.
-    # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
+#     # ℹ️ Command-line programs to run using the OS shell.
+#     # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
 
-    #   If the Autobuild fails above, remove it and uncomment the following three lines.
-    #   modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.
+#     #   If the Autobuild fails above, remove it and uncomment the following three lines.
+#     #   modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.
 
-    # - run: |
-    #   echo "Run, Build Application using script"
-    #   ./location_of_script_within_repo/buildscript.sh
+#     # - run: |
+#     #   echo "Run, Build Application using script"
+#     #   ./location_of_script_within_repo/buildscript.sh
 
-    - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v2
+#     - name: Perform CodeQL Analysis
+#       uses: github/codeql-action/analyze@v2

pyproject.toml

@@ -107,7 +107,7 @@ quote-style = "double"
 indent-style = "space"
 
 [tool.pyright]
-typeCheckingMode = "basic"
+typeCheckingMode = "strict"
 pythonVersion = "3.10"
 pythonPlatform = "All"
 

pytd/client.py

@@ -5,6 +5,9 @@
 
 import pandas as pd
 import tdclient
+from tdclient.models import Database as TDClientDatabase
+from tdclient.models import Job as TDClientJob
+from tdclient.models import Table as TDClientTable
 
 from .query_engine import HiveQueryEngine, PrestoQueryEngine, QueryEngine, QueryResult
 from .table import Table
@@ -119,7 +122,7 @@ def __init__(
             **kwargs,
         )
 
-    def list_databases(self) -> list[tdclient.models.Database]:
+    def list_databases(self) -> list[TDClientDatabase]:
         """Get a list of td-client-python Database objects.
 
         Returns
@@ -128,7 +131,7 @@ def list_databases(self) -> list[tdclient.models.Database]:
         """
         return self.api_client.databases()
 
-    def list_tables(self, database: str | None = None) -> list[tdclient.models.Table]:
+    def list_tables(self, database: str | None = None) -> list[TDClientTable]:
         """Get a list of td-client-python Table objects.
 
         Parameters
@@ -145,7 +148,7 @@ def list_tables(self, database: str | None = None) -> list[tdclient.models.Table
             database = self.database
         return self.api_client.tables(database)
 
-    def list_jobs(self) -> list[tdclient.models.Job]:
+    def list_jobs(self) -> list[TDClientJob]:
         """Get a list of td-client-python Job objects.
 
         Returns
@@ -154,7 +157,7 @@ def list_jobs(self) -> list[tdclient.models.Job]:
         """
         return self.api_client.jobs()
 
-    def get_job(self, job_id: int) -> tdclient.models.Job:
+    def get_job(self, job_id: int) -> TDClientJob:
         """Get a td-client-python Job object from ``job_id``.
 
         Parameters
@@ -370,7 +373,12 @@ def __exit__(
         self.close()
 
     def _fetch_query_engine(
-        self, engine: str, apikey: str, endpoint: str, database: str, header: str | bool
+        self,
+        engine: Literal["presto", "hive"],
+        apikey: str,
+        endpoint: str,
+        database: str,
+        header: str | bool,
     ) -> QueryEngine:
         if engine == "presto":
             return PrestoQueryEngine(apikey, endpoint, database, header)

pytd/pandas_td/__init__.py

@@ -199,7 +199,7 @@ def read_td_query(
         params = {}
     params["force_tdclient"] = True
 
-    if isinstance(engine, PrestoQueryEngine) and distributed_join is not None:
+    if isinstance(engine, PrestoQueryEngine):
         header = engine.create_header(
             [
                 "read_td_query",
@@ -337,7 +337,7 @@ def read_td_table(
     return _to_dataframe(engine.execute(query), index_col, parse_dates)
 
 
-def _convert_time(time):
+def _convert_time(time: Any) -> str:
     if time is None:
         return "NULL"
     elif isinstance(time, int):
@@ -351,7 +351,7 @@ def _convert_time(time):
     return f"'{t.replace(microsecond=0)}'"
 
 
-def _to_dataframe(dic, index_col, parse_dates):
+def _to_dataframe(dic: Any, index_col: Any, parse_dates: Any) -> pd.DataFrame:
     frame = pd.DataFrame(**dic)
     if parse_dates is not None:
         frame = _parse_dates(frame, parse_dates)
@@ -360,7 +360,7 @@ def _to_dataframe(dic, index_col, parse_dates):
     return frame
 
 
-def _parse_dates(frame, parse_dates):
+def _parse_dates(frame: pd.DataFrame, parse_dates: Any) -> pd.DataFrame:
     for name in parse_dates:
         if type(parse_dates) is list:
             frame[name] = pd.to_datetime(frame[name])
@@ -509,8 +509,6 @@ def _convert_time_column(
             frame["time"] = col.astype("int64") // (10**9)
     elif time_index is not None:
         # Use 'time_index' as time column
-        if type(time_index) is bool or not isinstance(time_index, int):
-            raise TypeError("invalid type for time_index")
         if isinstance(frame.index, pd.MultiIndex):
             idx = frame.index.levels[time_index]
         else:
@@ -533,8 +531,6 @@ def _convert_index_column(
     index: bool | None = None,
     index_label: str | list[str] | None = None,
 ) -> pd.DataFrame:
-    if index is not None and not isinstance(index, bool):
-        raise TypeError("index must be boolean")
     if index:
         if isinstance(frame.index, pd.MultiIndex):
             if index_label is None:
@@ -555,9 +551,9 @@ def _convert_date_format(
 ) -> pd.DataFrame:
     if date_format is not None:
 
-        def _convert(col):
+        def _convert(col: pd.Series) -> pd.Series:
             if col.dtype.name == "datetime64[ns]":
-                return col.apply(lambda x: x.strftime(date_format))
+                return col.apply(lambda x: x.strftime(date_format))  # type: ignore[return-value, misc]
             return col
 
         frame = frame.apply(_convert)  # type: ignore[assignment]

pytd/query_engine.py

@@ -6,8 +6,9 @@
 from urllib.parse import urlparse
 
 import tdclient
-import tdclient.cursor
 import trino
+from tdclient.connection import Connection as TDClientConnection
+from tdclient.cursor import Cursor as TDClientCursor
 
 __version__ = importlib.metadata.version("pytd")
 
@@ -45,7 +46,7 @@ def execute(
                 "params must be a list or tuple containing the query parameter values"
             )
 
-            if self.connection._use_legacy_prepared_statements():
+            if self.connection._use_legacy_prepared_statements():  # type: ignore[attr-defined]
                 statement_name = self._generate_unique_statement_name()
                 self._prepare_statement(operation, statement_name)
 
@@ -83,7 +84,7 @@ def execute(
 
 
 # Type alias for cursor types returned by query engines
-Cursor: TypeAlias = CustomTrinoCursor | tdclient.cursor.Cursor
+Cursor: TypeAlias = CustomTrinoCursor | TDClientCursor
 
 
 class QueryEngine(metaclass=abc.ABCMeta):
@@ -205,7 +206,7 @@ def create_header(
 
         if isinstance(extra_lines, str):
             header += f"-- {extra_lines}\n"
-        elif isinstance(extra_lines, list | tuple):
+        elif extra_lines:  # list or tuple
             header += "".join([f"-- {line}\n" for line in extra_lines])
 
         return header
@@ -223,8 +224,8 @@ def _connect(self) -> Any:
         pass
 
     def _get_tdclient_cursor(
-        self, con: tdclient.api.API, **kwargs: Any
-    ) -> tdclient.cursor.Cursor:
+        self, con: TDClientConnection, **kwargs: Any
+    ) -> TDClientCursor:
         """Get DB-API cursor from tdclient Connection instance.
 
         ``kwargs`` are for setting specific parameters to Treasure Data REST
@@ -290,7 +291,7 @@ def _get_tdclient_cursor(
             )
 
         # update a clone of the original params
-        cursor_kwargs = con._cursor_kwargs.copy()
+        cursor_kwargs = con._cursor_kwargs.copy()  # type: ignore[attr-defined]
         for k, v in kwargs.items():
             if k not in api_param_names:
                 raise RuntimeError(
@@ -300,17 +301,17 @@ def _get_tdclient_cursor(
             cursor_kwargs[k] = v
 
         # keep the original `_cursor_kwargs`
-        original_cursor_kwargs = con._cursor_kwargs.copy()
+        original_cursor_kwargs = con._cursor_kwargs.copy()  # type: ignore[attr-defined]
 
         # overwrite the original params
-        con._cursor_kwargs = cursor_kwargs
+        con._cursor_kwargs = cursor_kwargs  # type: ignore[attr-defined]
 
         # `Connection#cursor` internally refers the customized
         # ``_cursor_kwargs``
         cursor = con.cursor()
 
         # write the original params back to `_cursor_kwargs`
-        con._cursor_kwargs = original_cursor_kwargs
+        con._cursor_kwargs = original_cursor_kwargs  # type: ignore[attr-defined]
 
         logger.warning(
             "returning `tdclient.cursor.Cursor`. This cursor, `Cursor#fetchone` "
@@ -355,7 +356,7 @@ def user_agent(self) -> str:
         )
 
     @property
-    def presto_api_host(self):
+    def presto_api_host(self) -> str:
         """Presto API host obtained from ``TD_PRESTO_API`` env variable or
         inferred from Treasure Data REST API endpoint.
         """
@@ -402,7 +403,9 @@ def cursor(self, force_tdclient: bool = False, **kwargs: Any) -> Cursor:
 
             # In production, return our custom cursor with the same request object
             return CustomTrinoCursor(
-                self.trino_connection, regular_cursor._request, self.user_agent
+                self.trino_connection,
+                regular_cursor._request,  # type: ignore[attr-defined]
+                self.user_agent,  # type: ignore[attr-defined]
             )
 
         return self._get_tdclient_cursor(self.tdclient_connection, **kwargs)
@@ -414,7 +417,7 @@ def close(self) -> None:
 
     def _connect(
         self,
-    ) -> tuple[trino.dbapi.Connection, tdclient.api.API]:
+    ) -> tuple[trino.dbapi.Connection, TDClientConnection]:
         # Create trino connection
         trino_connection = trino.dbapi.connect(
             host=self.presto_api_host,
@@ -466,9 +469,7 @@ def user_agent(self) -> str:
         """User agent passed to a Hive connection."""
         return f"pytd/{__version__} (tdclient/{tdclient.__version__})"
 
-    def cursor(
-        self, force_tdclient: bool = True, **kwargs: Any
-    ) -> tdclient.cursor.Cursor:
+    def cursor(self, force_tdclient: bool = True, **kwargs: Any) -> TDClientCursor:
         """Get cursor defined by DB-API.
 
         Parameters
@@ -511,7 +512,7 @@ def close(self) -> None:
         """Close a connection to Hive."""
         self.engine.close()
 
-    def _connect(self) -> tdclient.api.API:
+    def _connect(self) -> TDClientConnection:
         return tdclient.connect(
             apikey=self.apikey,
             endpoint=self.endpoint,