Add type hints (cont.)

vertica_python/__init__.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .vertica.connection import Connection, connect, parse_dsn
 

vertica_python/datatypes.py

@@ -34,9 +34,12 @@
 # THE SOFTWARE.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from datetime import date, datetime, time, timezone
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from typing import Optional
 
 
 # noinspection PyPep8Naming
@@ -298,7 +301,7 @@ def __ne__(self, other):
 }
 
 def getTypeName(data_type_oid: int, type_modifier: int) -> str:
-    """Returns the base type name according to data_type_oid and type_modifier"""
+    """Returns the base type name according to data_type_oid and type_modifier."""
     if data_type_oid in TYPENAME:
         return TYPENAME[data_type_oid]
     elif data_type_oid in (VerticaType.INTERVAL, VerticaType.INTERVALYM):
@@ -310,12 +313,12 @@ def getTypeName(data_type_oid: int, type_modifier: int) -> str:
     else:
         return "Unknown"
 
-def getComplexElementType(data_type_oid: int):
-    """For 1D ARRAY or SET, returns the type of its elements"""
+def getComplexElementType(data_type_oid: int) -> Optional[int]:
+    """For 1D ARRAY or SET, returns the type of its elements."""
     return COMPLEX_ELEMENT_TYPE.get(data_type_oid)
 
-def getIntervalRange(data_type_oid: int, type_modifier: int):
-    """Extracts an interval's range from the bits set in its type_modifier"""
+def getIntervalRange(data_type_oid: int, type_modifier: int) -> str:
+    """Extracts an interval's range from the bits set in its type_modifier."""
 
     if data_type_oid not in (VerticaType.INTERVAL, VerticaType.INTERVALYM):
         raise ValueError("Invalid data type OID: {}".format(data_type_oid))
@@ -361,7 +364,7 @@ def getIntervalRange(data_type_oid: int, type_modifier: int):
             return "Day to Second"
 
 
-def getIntervalLeadingPrecision(data_type_oid: int, type_modifier: int):
+def getIntervalLeadingPrecision(data_type_oid: int, type_modifier: int) -> int:
     """
     Returns the leading precision for an interval, which is the largest number
     of digits that can fit in the leading field of the interval.
@@ -394,7 +397,7 @@ def getIntervalLeadingPrecision(data_type_oid: int, type_modifier: int):
         raise ValueError("Invalid interval range: {}".format(interval_range))
 
 
-def getPrecision(data_type_oid: int, type_modifier: int):
+def getPrecision(data_type_oid: int, type_modifier: int) -> Optional[int]:
     """
     Returns the precision for the given Vertica type with consideration of
     the type modifier.
@@ -423,22 +426,23 @@ def getPrecision(data_type_oid: int, type_modifier: int):
         return None  # None if no meaningful values can be provided
 
 
-def getScale(data_type_oid: int, type_modifier: int):
+def getScale(data_type_oid: int, type_modifier: int) -> Optional[int]:
     """
     Returns the scale for the given Vertica type with consideration of
-    the type modifier.
+    the type modifier. Returns None if no meaningful values can be provided.
     """
 
     if data_type_oid == VerticaType.NUMERIC:
         return 15 if type_modifier == -1 else (type_modifier - 4) & 0xFF
     else:
-        return None  # None if no meaningful values can be provided
+        return None
 
 
-def getDisplaySize(data_type_oid: int, type_modifier: int):
+def getDisplaySize(data_type_oid: int, type_modifier: int) -> Optional[int]:
     """
     Returns the column display size for the given Vertica type with
     consideration of the type modifier.
+    Returns None if no meaningful values can be provided.
 
     The display size of a column is the maximum number of characters needed to
     display data in character form.

vertica_python/errors.py

@@ -34,7 +34,7 @@
 # THE SOFTWARE.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import re
 

vertica_python/os_utils.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import errno
 import os

vertica_python/tests/common/base.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import os
 import sys

vertica_python/tests/integration_tests/base.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import pytest
 

vertica_python/tests/integration_tests/test_authentication.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/integration_tests/test_cancel.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from multiprocessing import Process
 import pytest

vertica_python/tests/integration_tests/test_column.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/integration_tests/test_connection.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import getpass
 import socket

vertica_python/tests/integration_tests/test_cursor.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from datetime import date, datetime, time
 from dateutil.relativedelta import relativedelta

vertica_python/tests/integration_tests/test_datatypes.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from datetime import date, datetime, time
 from dateutil.relativedelta import relativedelta

vertica_python/tests/integration_tests/test_dates.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from collections import namedtuple
 from datetime import date

vertica_python/tests/integration_tests/test_errors.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/integration_tests/test_loadbalance.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/integration_tests/test_timezones.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from collections import namedtuple
 from datetime import datetime

vertica_python/tests/integration_tests/test_tls.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import os
 import socket

vertica_python/tests/integration_tests/test_transfer_format.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/integration_tests/test_unicode.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
 

vertica_python/tests/unit_tests/base.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import pytest
 from ..common.base import VerticaPythonTestCase

vertica_python/tests/unit_tests/test_errors.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonUnitTestCase
 from ...errors import VerticaSyntaxError

vertica_python/tests/unit_tests/test_logging.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import logging
 import os

vertica_python/tests/unit_tests/test_notice.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 import mock
 

vertica_python/tests/unit_tests/test_parsedsn.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonUnitTestCase
 from ...vertica.connection import parse_dsn

vertica_python/tests/unit_tests/test_sql_literal.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from collections import namedtuple
 from decimal import Decimal

vertica_python/tests/unit_tests/test_timestamps.py

@@ -33,7 +33,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
 from collections import namedtuple
 from datetime import datetime

vertica_python/vertica/column.py

@@ -34,9 +34,11 @@
 # THE SOFTWARE.
 
 
-from __future__ import print_function, division, absolute_import
+from __future__ import print_function, division, absolute_import, annotations
 
-from collections import namedtuple
+from typing import TYPE_CHECKING, NamedTuple
+if TYPE_CHECKING:
+    from typing import Optional
 
 from ..datatypes import getDisplaySize, getPrecision, getScale
 from ..compat import as_str, as_text
@@ -49,12 +51,18 @@ class FormatCode(object):
     BINARY = 1
 
 
-ColumnTuple = namedtuple('Column', ['name', 'type_code', 'display_size', 'internal_size',
-                                    'precision', 'scale', 'null_ok'])
+class ColumnTuple(NamedTuple):
+    name: str
+    type_code: int
+    display_size: Optional[int]
+    internal_size: int
+    precision: Optional[int]
+    scale: Optional[int]
+    null_ok: bool
 
 
 class Column(object):
-    def __init__(self, col):
+    def __init__(self, col) -> None:
         # Describe one query result column
         self.name = col['name']
         self.type_code = col['data_type_oid']
@@ -74,7 +82,7 @@ def __init__(self, col):
         self.props = ColumnTuple(self.name, self.type_code, self.display_size, self.internal_size,
                                  self.precision, self.scale, self.null_ok)
 
-    def add_child_column(self, col):
+    def add_child_column(self, col: Column) -> None:
         """
         Complex types involve multiple columns arranged in a hierarchy of parents and children.
         Each parent column stores references to child columns in a list.
@@ -83,7 +91,7 @@ def add_child_column(self, col):
             self.child_columns = []
         self.child_columns.append(col)
 
-    def debug_info(self):
+    def debug_info(self) -> str:
         childs = ""
         if self.child_columns:
             c = ", ".join([col.debug_info() for col in self.child_columns])

vertica_python/vertica/connection.py

@@ -633,7 +633,7 @@ def ssl(self) -> bool:
         """Returns True if the TCP socket is a SSL socket."""
         return self.socket is not None and isinstance(self.socket, ssl.SSLSocket)
 
-    def write(self, message, vsocket=None) -> None:
+    def write(self, message: FrontendMessage, vsocket: Optional[Union[socket.socket, ssl.SSLSocket]] = None) -> None:
         if not isinstance(message, FrontendMessage):
             raise TypeError("invalid message: ({0})".format(message))
         if vsocket is None:
@@ -670,14 +670,14 @@ def reset_connection(self) -> None:
         self.close()
         self.startup_connection()
 
-    def is_asynchronous_message(self, message) -> bool:
+    def is_asynchronous_message(self, message: BackendMessage) -> bool:
         # Check if it is an asynchronous response message
         # Note: ErrorResponse is a subclass of NoticeResponse
         return (isinstance(message, messages.ParameterStatus) or
             (isinstance(message, messages.NoticeResponse) and
              not isinstance(message, messages.ErrorResponse)))
 
-    def handle_asynchronous_message(self, message) -> None:
+    def handle_asynchronous_message(self, message: Union[messages.ParameterStatus, messages.NoticeResponse]) -> None:
         if isinstance(message, messages.ParameterStatus):
             if message.name == 'protocol_version':
                 message.value = int(message.value)
@@ -700,7 +700,7 @@ def read_string(self) -> bytearray:
             s.extend(char)
         return s
 
-    def read_message(self):
+    def read_message(self) -> BackendMessage:
         while True:
             try:
                 type_ = self.read_bytes(1)