singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/converters.py

@@ -0,0 +1,950 @@
+#!/usr/bin/env python
+"""Data value conversion utilities."""
+import datetime
+import re
+import struct
+from base64 import b64decode
+from decimal import Decimal
+from json import loads as json_loads
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Set
+from typing import Union
+
+try:
+    import shapely.wkt
+    has_shapely = True
+except ImportError:
+    has_shapely = False
+
+try:
+    import pygeos
+    has_pygeos = True
+except (AttributeError, ImportError):
+    has_pygeos = False
+
+try:
+    import numpy
+    has_numpy = True
+except ImportError:
+    has_numpy = False
+
+try:
+    import bson
+    has_bson = True
+except ImportError:
+    has_bson = False
+
+
+# Cache fromisoformat methods if they exist
+# Cache fromisoformat methods if they exist
+_dt_datetime_fromisoformat = None
+if hasattr(datetime.datetime, 'fromisoformat'):
+    _dt_datetime_fromisoformat = datetime.datetime.fromisoformat  # type: ignore
+_dt_time_fromisoformat = None
+if hasattr(datetime.time, 'fromisoformat'):
+    _dt_time_fromisoformat = datetime.time.fromisoformat  # type: ignore
+_dt_date_fromisoformat = None
+if hasattr(datetime.date, 'fromisoformat'):
+    _dt_date_fromisoformat = datetime.date.fromisoformat  # type: ignore
+
+
+def _convert_second_fraction(s: str) -> int:
+    if not s:
+        return 0
+    # Pad zeros to ensure the fraction length in microseconds
+    s = s.ljust(6, '0')
+    return int(s[:6])
+
+
+DATETIME_RE = re.compile(
+    r'(\d{1,4})-(\d{1,2})-(\d{1,2})[T ](\d{1,2}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?',
+)
+
+ZERO_DATETIMES = set([
+    '0000-00-00 00:00:00',
+    '0000-00-00 00:00:00.000',
+    '0000-00-00 00:00:00.000000',
+    '0000-00-00T00:00:00',
+    '0000-00-00T00:00:00.000',
+    '0000-00-00T00:00:00.000000',
+])
+ZERO_DATES = set([
+    '0000-00-00',
+])
+
+
+def datetime_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.datetime, str, None]:
+    """Returns a DATETIME or TIMESTAMP column value as a datetime object:
+
+      >>> datetime_fromisoformat('2007-02-25 23:06:20')
+      datetime.datetime(2007, 2, 25, 23, 6, 20)
+      >>> datetime_fromisoformat('2007-02-25T23:06:20')
+      datetime.datetime(2007, 2, 25, 23, 6, 20)
+
+    Illegal values are returned as str or None:
+
+      >>> datetime_fromisoformat('2007-02-31T23:06:20')
+      '2007-02-31T23:06:20'
+      >>> datetime_fromisoformat('0000-00-00 00:00:00')
+      None
+
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    if obj in ZERO_DATETIMES:
+        return None
+
+    # Use datetime methods if possible
+    if _dt_datetime_fromisoformat is not None:
+        try:
+            if ' ' in obj or 'T' in obj:
+                return _dt_datetime_fromisoformat(obj)
+            if _dt_date_fromisoformat is not None:
+                date = _dt_date_fromisoformat(obj)
+                return datetime.datetime(date.year, date.month, date.day)
+        except ValueError:
+            return obj
+
+    m = DATETIME_RE.match(obj)
+    if not m:
+        mdate = date_fromisoformat(obj)
+        if type(mdate) is str:
+            return mdate
+        return datetime.datetime(mdate.year, mdate.month, mdate.day)  # type: ignore
+
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        return datetime.datetime(*[int(x) for x in groups])  # type: ignore
+    except ValueError:
+        mdate = date_fromisoformat(obj)
+        if type(mdate) is str:
+            return mdate
+        return datetime.datetime(mdate.year, mdate.month, mdate.day)  # type: ignore
+
+
+TIMEDELTA_RE = re.compile(r'(-)?(\d{1,3}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?')
+
+
+def timedelta_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.timedelta, str, None]:
+    """Returns a TIME column as a timedelta object:
+
+      >>> timedelta_fromisoformat('25:06:17')
+      datetime.timedelta(days=1, seconds=3977)
+      >>> timedelta_fromisoformat('-25:06:17')
+      datetime.timedelta(days=-2, seconds=82423)
+
+    Illegal values are returned as string:
+
+      >>> timedelta_fromisoformat('random crap')
+      'random crap'
+
+    Note that MySQL always returns TIME columns as (+|-)HH:MM:SS, but
+    can accept values as (+|-)DD HH:MM:SS. The latter format will not
+    be parsed correctly by this function.
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    m = TIMEDELTA_RE.match(obj)
+    if not m:
+        return obj
+
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        negate = -1 if groups[0] else 1
+        hours, minutes, seconds, microseconds = groups[1:]
+
+        tdelta = (
+            datetime.timedelta(
+                hours=int(hours),
+                minutes=int(minutes),
+                seconds=int(seconds),
+                microseconds=int(microseconds),
+            )
+            * negate
+        )
+        return tdelta
+    except ValueError:
+        return obj
+
+
+TIME_RE = re.compile(r'(\d{1,2}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?')
+
+
+def time_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.time, str, None]:
+    """Returns a TIME column as a time object:
+
+      >>> time_fromisoformat('15:06:17')
+      datetime.time(15, 6, 17)
+
+    Illegal values are returned as str:
+
+      >>> time_fromisoformat('-25:06:17')
+      '-25:06:17'
+      >>> time_fromisoformat('random crap')
+      'random crap'
+
+    Note that MySQL always returns TIME columns as (+|-)HH:MM:SS, but
+    can accept values as (+|-)DD HH:MM:SS. The latter format will not
+    be parsed correctly by this function.
+
+    Also note that MySQL's TIME column corresponds more closely to
+    Python's timedelta and not time. However if you want TIME columns
+    to be treated as time-of-day and not a time offset, then you can
+    use set this function as the converter for FIELD_TYPE.TIME.
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    # Use datetime methods if possible
+    if _dt_time_fromisoformat is not None:
+        try:
+            return _dt_time_fromisoformat(obj)
+        except ValueError:
+            return obj
+
+    m = TIME_RE.match(obj)
+    if not m:
+        return obj
+
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        hours, minutes, seconds, microseconds = groups
+        return datetime.time(
+            hour=int(hours),
+            minute=int(minutes),
+            second=int(seconds),
+            microsecond=int(microseconds),
+        )
+    except ValueError:
+        return obj
+
+
+def date_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.date, str, None]:
+    """Returns a DATE column as a date object:
+
+      >>> date_fromisoformat('2007-02-26')
+      datetime.date(2007, 2, 26)
+
+    Illegal values are returned as str or None:
+
+      >>> date_fromisoformat('2007-02-31')
+      '2007-02-31'
+      >>> date_fromisoformat('0000-00-00')
+      None
+
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    if obj in ZERO_DATES:
+        return None
+
+    # Use datetime methods if possible
+    if _dt_date_fromisoformat is not None:
+        try:
+            return _dt_date_fromisoformat(obj)
+        except ValueError:
+            return obj
+
+    try:
+        return datetime.date(*[int(x) for x in obj.split('-', 2)])
+    except ValueError:
+        return obj
+
+
+def identity(x: Any) -> Optional[Any]:
+    """Return input value."""
+    return x
+
+
+def bit_or_none(x: Any) -> Optional[bytes]:
+    """
+    Convert value to bit.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    int
+        If value can be cast to a bit
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    if type(x) is str:
+        return b64decode(x)
+    return x
+
+
+def int_or_none(x: Any) -> Optional[int]:
+    """
+    Convert value to int.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    int
+        If value can be cast to an int
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return int(x)
+
+
+def float_or_none(x: Any) -> Optional[float]:
+    """
+    Convert value to float.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    float
+        If value can be cast to a float
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return float(x)
+
+
+def decimal_or_none(x: Any) -> Optional[Decimal]:
+    """
+    Convert value to decimal.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    decimal.Decimal
+        If value can be cast to a decimal
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return Decimal(x)
+
+
+def date_or_none(x: Optional[str]) -> Optional[Union[datetime.date, str]]:
+    """
+    Convert value to a date.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    datetime.date
+        If value can be cast to a date
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return date_fromisoformat(x)
+
+
+def timedelta_or_none(x: Optional[str]) -> Optional[Union[datetime.timedelta, str]]:
+    """
+    Convert value to a timedelta.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    datetime.timedelta
+        If value can be cast to a time
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return timedelta_fromisoformat(x)
+
+
+def time_or_none(x: Optional[str]) -> Optional[Union[datetime.time, str]]:
+    """
+    Convert value to a time.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    datetime.time
+        If value can be cast to a time
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return time_fromisoformat(x)
+
+
+def datetime_or_none(x: Optional[str]) -> Optional[Union[datetime.datetime, str]]:
+    """
+    Convert value to a datetime.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    datetime.time
+        If value can be cast to a datetime
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return datetime_fromisoformat(x)
+
+
+def none(x: Any) -> None:
+    """
+    Return None.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    None
+
+    """
+    return None
+
+
+def json_or_none(x: Optional[str]) -> Optional[Union[Dict[str, Any], List[Any]]]:
+    """
+    Convert JSON to dict or list.
+
+    Parameters
+    ----------
+    x : str
+        JSON string
+
+    Returns
+    -------
+    dict
+        If JSON string contains an object
+    list
+        If JSON string contains a list
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return json_loads(x)
+
+
+def set_or_none(x: Optional[str]) -> Optional[Set[str]]:
+    """
+    Convert value to set of strings.
+
+    Parameters
+    ----------
+    x : str
+        Input string value
+
+    Returns
+    -------
+    set of strings
+        If value contains a set of strings
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    return set(y.strip() for y in x.split(','))
+
+
+def geometry_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Convert value to geometry coordinates.
+
+    Parameters
+    ----------
+    x : Any
+        Geometry value
+
+    Returns
+    -------
+    shapely object or pygeos object or str
+        If value is valid geometry value
+    None
+        If input value is None or empty
+
+    """
+    if x is None or not x:
+        return None
+    if has_shapely:
+        return shapely.wkt.loads(x)
+    if has_pygeos:
+        return pygeos.io.from_wkt(x)
+    return x
+
+
+def float32_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to float32 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    float32 numpy array
+        If input value is not None and numpy is installed
+    float Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.float32)
+
+    return map(float, json_loads(x))
+
+
+def float32_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to float32 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        Little-endian block of bytes.
+
+    Returns
+    -------
+    float32 numpy array
+        If input value is not None and numpy is installed
+    float Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.float32)
+
+    return struct.unpack(f'<{len(x)//4}f', x)
+
+
+def float64_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to float64 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    float64 numpy array
+        If input value is not None and numpy is installed
+    float Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.float64)
+
+    return map(float, json_loads(x))
+
+
+def float64_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to float64 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        JSON array
+
+    Returns
+    -------
+    float64 numpy array
+        If input value is not None and numpy is installed
+    float Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.float64)
+
+    return struct.unpack(f'<{len(x)//8}d', x)
+
+
+def int8_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to int8 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    int8 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.int8)
+
+    return map(int, json_loads(x))
+
+
+def int8_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to int8 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        Little-endian block of bytes.
+
+    Returns
+    -------
+    int8 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.int8)
+
+    return struct.unpack(f'<{len(x)}b', x)
+
+
+def int16_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to int16 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    int16 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.int16)
+
+    return map(int, json_loads(x))
+
+
+def int16_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to int16 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        Little-endian block of bytes.
+
+    Returns
+    -------
+    int16 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.int16)
+
+    return struct.unpack(f'<{len(x)//2}h', x)
+
+
+def int32_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to int32 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    int32 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.int32)
+
+    return map(int, json_loads(x))
+
+
+def int32_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to int32 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        Little-endian block of bytes.
+
+    Returns
+    -------
+    int32 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.int32)
+
+    return struct.unpack(f'<{len(x)//4}l', x)
+
+
+def int64_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
+    """
+    Covert value to int64 array.
+
+    Parameters
+    ----------
+    x : str or None
+        JSON array
+
+    Returns
+    -------
+    int64 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    if has_numpy:
+        return numpy.array(json_loads(x), dtype=numpy.int64)
+
+    return map(int, json_loads(x))
+
+
+def int64_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Covert value to int64 array.
+
+    Parameters
+    ----------
+    x : bytes or None
+        Little-endian block of bytes.
+
+    Returns
+    -------
+    int64 numpy array
+        If input value is not None and numpy is installed
+    int Python list
+        If input value is not None and numpy is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+
+    # Bytes
+    if has_numpy:
+        return numpy.frombuffer(x, dtype=numpy.int64)
+
+    return struct.unpack(f'<{len(x)//8}l', x)
+
+
+def bson_or_none(x: Optional[bytes]) -> Optional[Any]:
+    """
+    Convert a BSON value to a dictionary.
+
+    Parameters
+    ----------
+    x : bytes or None
+        BSON formatted bytes
+
+    Returns
+    -------
+    dict
+        If input value is not None and bson package is installed
+    bytes
+        If input value is not None and bson package is not installed
+    None
+        If input value is None
+
+    """
+    if x is None:
+        return None
+    if has_bson:
+        return bson.decode(x)
+    return x
+
+
+# Map of database types and conversion functions
+converters: Dict[int, Callable[..., Any]] = {
+    0: decimal_or_none,
+    1: int_or_none,
+    2: int_or_none,
+    3: int_or_none,
+    4: float_or_none,
+    5: float_or_none,
+    6: none,
+    7: datetime_or_none,
+    8: int_or_none,
+    9: int_or_none,
+    10: date_or_none,
+    11: timedelta_or_none,
+    12: datetime_or_none,
+    13: int_or_none,
+    14: date_or_none,
+    #   15: identity,
+    16: bit_or_none,
+    245: json_or_none,
+    246: decimal_or_none,
+    #   247: identity,
+    248: set_or_none,
+    #   249: identity,
+    #   250: identity,
+    #   251: identity,
+    #   252: identity,
+    #   253: identity,
+    #   254: identity,
+    255: geometry_or_none,
+    1001: bson_or_none,
+    2001: float32_vector_json_or_none,
+    2002: float64_vector_json_or_none,
+    2003: int8_vector_json_or_none,
+    2004: int16_vector_json_or_none,
+    2005: int32_vector_json_or_none,
+    2006: int64_vector_json_or_none,
+    3001: float32_vector_or_none,
+    3002: float64_vector_or_none,
+    3003: int8_vector_or_none,
+    3004: int16_vector_or_none,
+    3005: int32_vector_or_none,
+    3006: int64_vector_or_none,
+}