meerschaum.utils.dtypes

  1#! /usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Utility functions for working with data types.
  7"""
  8
  9import traceback
 10import json
 11import uuid
 12from datetime import timezone, datetime
 13from decimal import Decimal, Context, InvalidOperation, ROUND_HALF_UP
 14
 15import meerschaum as mrsm
 16from meerschaum.utils.typing import Dict, Union, Any, Optional, Tuple
 17from meerschaum.utils.warnings import warn
 18
 19MRSM_ALIAS_DTYPES: Dict[str, str] = {
 20    'decimal': 'numeric',
 21    'Decimal': 'numeric',
 22    'number': 'numeric',
 23    'jsonl': 'json',
 24    'JSON': 'json',
 25    'binary': 'bytes',
 26    'blob': 'bytes',
 27    'varbinary': 'bytes',
 28    'bytea': 'bytes',
 29    'guid': 'uuid',
 30    'UUID': 'uuid',
 31    'geom': 'geometry',
 32    'geog': 'geography',
 33}
 34MRSM_PD_DTYPES: Dict[Union[str, None], str] = {
 35    'json': 'object',
 36    'numeric': 'object',
 37    'geometry': 'object',
 38    'geography': 'object',
 39    'uuid': 'object',
 40    'datetime': 'datetime64[ns, UTC]',
 41    'bool': 'bool[pyarrow]',
 42    'int': 'Int64',
 43    'int8': 'Int8',
 44    'int16': 'Int16',
 45    'int32': 'Int32',
 46    'int64': 'Int64',
 47    'str': 'string[python]',
 48    'bytes': 'object',
 49    None: 'object',
 50}
 51
 52
 53def to_pandas_dtype(dtype: str) -> str:
 54    """
 55    Cast a supported Meerschaum dtype to a Pandas dtype.
 56    """
 57    known_dtype = MRSM_PD_DTYPES.get(dtype, None)
 58    if known_dtype is not None:
 59        return known_dtype
 60
 61    alias_dtype = MRSM_ALIAS_DTYPES.get(dtype, None)
 62    if alias_dtype is not None:
 63        return MRSM_PD_DTYPES[alias_dtype]
 64
 65    if dtype.startswith('numeric'):
 66        return MRSM_PD_DTYPES['numeric']
 67
 68    if dtype.startswith('geometry'):
 69        return MRSM_PD_DTYPES['geometry']
 70
 71    if dtype.startswith('geography'):
 72        return MRSM_PD_DTYPES['geography']
 73
 74    ### NOTE: Kind of a hack, but if the first word of the given dtype is in all caps,
 75    ### treat it as a SQL db type.
 76    if dtype.split(' ')[0].isupper():
 77        from meerschaum.utils.dtypes.sql import get_pd_type_from_db_type
 78        return get_pd_type_from_db_type(dtype)
 79
 80    from meerschaum.utils.packages import attempt_import
 81    _ = attempt_import('pyarrow', lazy=False)
 82    pandas = attempt_import('pandas', lazy=False)
 83
 84    try:
 85        return str(pandas.api.types.pandas_dtype(dtype))
 86    except Exception:
 87        warn(
 88            f"Invalid dtype '{dtype}', will use 'object' instead:\n"
 89            + f"{traceback.format_exc()}",
 90            stack=False,
 91        )
 92    return 'object'
 93
 94
 95def are_dtypes_equal(
 96    ldtype: Union[str, Dict[str, str]],
 97    rdtype: Union[str, Dict[str, str]],
 98) -> bool:
 99    """
100    Determine whether two dtype strings may be considered
101    equivalent to avoid unnecessary conversions.
102
103    Parameters
104    ----------
105    ldtype: Union[str, Dict[str, str]]
106        The left dtype to compare.
107        May also provide a dtypes dictionary.
108
109    rdtype: Union[str, Dict[str, str]]
110        The right dtype to compare.
111        May also provide a dtypes dictionary.
112
113    Returns
114    -------
115    A `bool` indicating whether the two dtypes are to be considered equivalent.
116    """
117    if isinstance(ldtype, dict) and isinstance(rdtype, dict):
118        lkeys = sorted([str(k) for k in ldtype.keys()])
119        rkeys = sorted([str(k) for k in rdtype.keys()])
120        for lkey, rkey in zip(lkeys, rkeys):
121            if lkey != rkey:
122                return False
123            ltype = ldtype[lkey]
124            rtype = rdtype[rkey]
125            if not are_dtypes_equal(ltype, rtype):
126                return False
127        return True
128
129    try:
130        if ldtype == rdtype:
131            return True
132    except Exception:
133        warn(f"Exception when comparing dtypes, returning False:\n{traceback.format_exc()}")
134        return False
135
136    ### Sometimes pandas dtype objects are passed.
137    ldtype = str(ldtype).split('[', maxsplit=1)[0]
138    rdtype = str(rdtype).split('[', maxsplit=1)[0]
139
140    if ldtype in MRSM_ALIAS_DTYPES:
141        ldtype = MRSM_ALIAS_DTYPES[ldtype]
142
143    if rdtype in MRSM_ALIAS_DTYPES:
144        rdtype = MRSM_ALIAS_DTYPES[rdtype]
145
146    json_dtypes = ('json', 'object')
147    if ldtype in json_dtypes and rdtype in json_dtypes:
148        return True
149
150    numeric_dtypes = ('numeric', 'object')
151    if ldtype in numeric_dtypes and rdtype in numeric_dtypes:
152        return True
153
154    uuid_dtypes = ('uuid', 'object')
155    if ldtype in uuid_dtypes and rdtype in uuid_dtypes:
156        return True
157
158    bytes_dtypes = ('bytes', 'object')
159    if ldtype in bytes_dtypes and rdtype in bytes_dtypes:
160        return True
161
162    geometry_dtypes = ('geometry', 'object', 'geography')
163    if ldtype in geometry_dtypes and rdtype in geometry_dtypes:
164        return True
165
166    if ldtype.lower() == rdtype.lower():
167        return True
168
169    datetime_dtypes = ('datetime', 'timestamp')
170    ldtype_found_dt_prefix = False
171    rdtype_found_dt_prefix = False
172    for dt_prefix in datetime_dtypes:
173        ldtype_found_dt_prefix = (dt_prefix in ldtype.lower()) or ldtype_found_dt_prefix
174        rdtype_found_dt_prefix = (dt_prefix in rdtype.lower()) or rdtype_found_dt_prefix
175    if ldtype_found_dt_prefix and rdtype_found_dt_prefix:
176        return True
177
178    string_dtypes = ('str', 'string', 'object')
179    if ldtype in string_dtypes and rdtype in string_dtypes:
180        return True
181
182    int_dtypes = ('int', 'int64', 'int32', 'int16', 'int8')
183    if ldtype.lower() in int_dtypes and rdtype.lower() in int_dtypes:
184        return True
185
186    float_dtypes = ('float', 'float64', 'float32', 'float16', 'float128', 'double')
187    if ldtype.lower() in float_dtypes and rdtype.lower() in float_dtypes:
188        return True
189
190    bool_dtypes = ('bool', 'boolean')
191    if ldtype in bool_dtypes and rdtype in bool_dtypes:
192        return True
193
194    return False
195
196
197def is_dtype_numeric(dtype: str) -> bool:
198    """
199    Determine whether a given `dtype` string
200    should be considered compatible with the Meerschaum dtype `numeric`.
201
202    Parameters
203    ----------
204    dtype: str
205        The pandas-like dtype string.
206
207    Returns
208    -------
209    A bool indicating the dtype is compatible with `numeric`.
210    """
211    dtype_lower = dtype.lower()
212
213    acceptable_substrings = ('numeric', 'float', 'double', 'int')
214    for substring in acceptable_substrings:
215        if substring in dtype_lower:
216            return True
217
218    return False
219
220
221def attempt_cast_to_numeric(
222    value: Any,
223    quantize: bool = False,
224    precision: Optional[int] = None,
225    scale: Optional[int] = None,
226)-> Any:
227    """
228    Given a value, attempt to coerce it into a numeric (Decimal).
229
230    Parameters
231    ----------
232    value: Any
233        The value to be cast to a Decimal.
234
235    quantize: bool, default False
236        If `True`, quantize the decimal to the specified precision and scale.
237
238    precision: Optional[int], default None
239        If `quantize` is `True`, use this precision.
240
241    scale: Optional[int], default None
242        If `quantize` is `True`, use this scale.
243
244    Returns
245    -------
246    A `Decimal` if possible, or `value`.
247    """
248    if isinstance(value, Decimal):
249        if quantize and precision and scale:
250            return quantize_decimal(value, precision, scale)
251        return value
252    try:
253        if value_is_null(value):
254            return Decimal('NaN')
255
256        dec = Decimal(str(value))
257        if not quantize or not precision or not scale:
258            return dec
259        return quantize_decimal(dec, precision, scale)
260    except Exception:
261        return value
262
263
264def attempt_cast_to_uuid(value: Any) -> Any:
265    """
266    Given a value, attempt to coerce it into a UUID (`uuid4`).
267    """
268    if isinstance(value, uuid.UUID):
269        return value
270    try:
271        return (
272            uuid.UUID(str(value))
273            if not value_is_null(value)
274            else None
275        )
276    except Exception:
277        return value
278
279
280def attempt_cast_to_bytes(value: Any) -> Any:
281    """
282    Given a value, attempt to coerce it into a bytestring.
283    """
284    if isinstance(value, bytes):
285        return value
286    try:
287        return (
288            deserialize_bytes_string(str(value))
289            if not value_is_null(value)
290            else None
291        )
292    except Exception:
293        return value
294
295
296def attempt_cast_to_geometry(value: Any) -> Any:
297    """
298    Given a value, attempt to coerce it into a `shapely` (`geometry`) object.
299    """
300    shapely, shapely_wkt, shapely_wkb = mrsm.attempt_import(
301        'shapely',
302        'shapely.wkt',
303        'shapely.wkb',
304        lazy=False,
305    )
306    if 'shapely' in str(type(value)):
307        return value
308
309    if isinstance(value, (dict, list)):
310        try:
311            return shapely.from_geojson(json.dumps(value))
312        except Exception as e:
313            return value
314
315    value_is_wkt = geometry_is_wkt(value)
316    if value_is_wkt is None:
317        return value
318
319    try:
320        return (
321            shapely_wkt.loads(value)
322            if value_is_wkt
323            else shapely_wkb.loads(value)
324        )
325    except Exception:
326        return value
327
328
329def geometry_is_wkt(value: Union[str, bytes]) -> Union[bool, None]:
330    """
331    Determine whether an input value should be treated as WKT or WKB geometry data.
332
333    Parameters
334    ----------
335    value: Union[str, bytes]
336        The input data to be parsed into geometry data.
337
338    Returns
339    -------
340    A `bool` (`True` if `value` is WKT and `False` if it should be treated as WKB).
341    Return `None` if `value` should be parsed as neither.
342    """
343    import re
344    if not isinstance(value, (str, bytes)):
345        return None
346
347    if isinstance(value, bytes):
348        return False
349    
350    wkt_pattern = r'^\s*(POINT|LINESTRING|POLYGON|MULTIPOINT|MULTILINESTRING|MULTIPOLYGON|GEOMETRYCOLLECTION)\s*\(.*\)\s*$'
351    if re.match(wkt_pattern, value, re.IGNORECASE):
352        return True
353    
354    if all(c in '0123456789ABCDEFabcdef' for c in value) and len(value) % 2 == 0:
355        return False
356    
357    return None
358
359
360def value_is_null(value: Any) -> bool:
361    """
362    Determine if a value is a null-like string.
363    """
364    return str(value).lower() in ('none', 'nan', 'na', 'nat', '', '<na>')
365
366
367def none_if_null(value: Any) -> Any:
368    """
369    Return `None` if a value is a null-like string.
370    """
371    return (None if value_is_null(value) else value)
372
373
374def quantize_decimal(x: Decimal, precision: int, scale: int) -> Decimal:
375    """
376    Quantize a given `Decimal` to a known scale and precision.
377
378    Parameters
379    ----------
380    x: Decimal
381        The `Decimal` to be quantized.
382
383    precision: int
384        The total number of significant digits.
385
386    scale: int
387        The number of significant digits after the decimal point.
388
389    Returns
390    -------
391    A `Decimal` quantized to the specified scale and precision.
392    """
393    precision_decimal = Decimal(('1' * (precision - scale)) + '.' + ('1' * scale))
394    try:
395        return x.quantize(precision_decimal, context=Context(prec=precision), rounding=ROUND_HALF_UP)
396    except InvalidOperation:
397        pass
398
399    raise ValueError(f"Cannot quantize value '{x}' to {precision=}, {scale=}.")
400
401
402def serialize_decimal(
403    x: Any,
404    quantize: bool = False,
405    precision: Optional[int] = None,
406    scale: Optional[int] = None,
407) -> Any:
408    """
409    Return a quantized string of an input decimal.
410
411    Parameters
412    ----------
413    x: Any
414        The potential decimal to be serialized.
415
416    quantize: bool, default False
417        If `True`, quantize the incoming Decimal to the specified scale and precision
418        before serialization.
419
420    precision: Optional[int], default None
421        The precision of the decimal to be quantized.
422
423    scale: Optional[int], default None
424        The scale of the decimal to be quantized.
425
426    Returns
427    -------
428    A string of the input decimal or the input if not a Decimal.
429    """
430    if not isinstance(x, Decimal):
431        return x
432
433    if value_is_null(x):
434        return None
435
436    if quantize and scale and precision:
437        x = quantize_decimal(x, precision, scale)
438
439    return f"{x:f}"
440
441
442def coerce_timezone(
443    dt: Any,
444    strip_utc: bool = False,
445) -> Any:
446    """
447    Given a `datetime`, pandas `Timestamp` or `Series` of `Timestamp`,
448    return a UTC timestamp (strip timezone if `strip_utc` is `True`.
449    """
450    if dt is None:
451        return None
452
453    if isinstance(dt, int):
454        return dt
455
456    if isinstance(dt, str):
457        dateutil_parser = mrsm.attempt_import('dateutil.parser')
458        dt = dateutil_parser.parse(dt)
459
460    dt_is_series = hasattr(dt, 'dtype') and hasattr(dt, '__module__')
461
462    if dt_is_series:
463        pandas = mrsm.attempt_import('pandas', lazy=False)
464
465        if (
466            pandas.api.types.is_datetime64_any_dtype(dt) and (
467                (dt.dt.tz is not None and not strip_utc)
468                or
469                (dt.dt.tz is None and strip_utc)
470            )
471        ):
472            return dt
473
474        dt_series = to_datetime(dt, coerce_utc=False)
475        if strip_utc:
476            try:
477                if dt_series.dt.tz is not None:
478                    dt_series = dt_series.dt.tz_localize(None)
479            except Exception:
480                pass
481
482        return dt_series
483
484    if dt.tzinfo is None:
485        if strip_utc:
486            return dt
487        return dt.replace(tzinfo=timezone.utc)
488
489    utc_dt = dt.astimezone(timezone.utc)
490    if strip_utc:
491        return utc_dt.replace(tzinfo=None)
492    return utc_dt
493
494
495def to_datetime(dt_val: Any, as_pydatetime: bool = False, coerce_utc: bool = True) -> Any:
496    """
497    Wrap `pd.to_datetime()` and add support for out-of-bounds values.
498    """
499    pandas, dateutil_parser = mrsm.attempt_import('pandas', 'dateutil.parser', lazy=False)
500    is_dask = 'dask' in getattr(dt_val, '__module__', '')
501    dd = mrsm.attempt_import('dask.dataframe') if is_dask else None
502    dt_is_series = hasattr(dt_val, 'dtype') and hasattr(dt_val, '__module__')
503    pd = pandas if dd is None else dd
504
505    try:
506        new_dt_val = pd.to_datetime(dt_val, utc=True, format='ISO8601')
507        if as_pydatetime:
508            return new_dt_val.to_pydatetime()
509        return new_dt_val
510    except (pd.errors.OutOfBoundsDatetime, ValueError):
511        pass
512
513    def parse(x: Any) -> Any:
514        try:
515            return dateutil_parser.parse(x)
516        except Exception:
517            return x
518
519    if dt_is_series:
520        new_series = dt_val.apply(parse)
521        if coerce_utc:
522            return coerce_timezone(new_series)
523        return new_series
524
525    new_dt_val = parse(dt_val)
526    if not coerce_utc:
527        return new_dt_val
528    return coerce_timezone(new_dt_val)
529
530
531def serialize_bytes(data: bytes) -> str:
532    """
533    Return the given bytes as a base64-encoded string.
534    """
535    import base64
536    if not isinstance(data, bytes) and value_is_null(data):
537        return data
538    return base64.b64encode(data).decode('utf-8')
539
540
541def serialize_geometry(
542    geom: Any,
543    geometry_format: str = 'wkb_hex',
544    as_wkt: bool = False,
545) -> Union[str, Dict[str, Any], None]:
546    """
547    Serialize geometry data as a hex-encoded well-known-binary string. 
548
549    Parameters
550    ----------
551    geom: Any
552        The potential geometry data to be serialized.
553
554    geometry_format: str, default 'wkb_hex'
555        The serialization format for geometry data.
556        Accepted formats are `wkb_hex` (well-known binary hex string),
557        `wkt` (well-known text), and `geojson`.
558
559    Returns
560    -------
561    A string containing the geometry data.
562    """
563    if value_is_null(geom):
564        return None
565    shapely = mrsm.attempt_import('shapely', lazy=False)
566    if geometry_format == 'geojson':
567        geojson_str = shapely.to_geojson(geom)
568        return json.loads(geojson_str)
569
570    if hasattr(geom, 'wkb_hex'):
571        return geom.wkb_hex if geometry_format == 'wkb_hex' else geom.wkt
572
573    return str(geom)
574
575
576def deserialize_geometry(geom_wkb: Union[str, bytes]):
577    """
578    Deserialize a WKB string into a shapely geometry object.
579    """
580    shapely = mrsm.attempt_import(lazy=False)
581    return shapely.wkb.loads(geom_wkb)
582
583
584def deserialize_bytes_string(data: Optional[str], force_hex: bool = False) -> Union[bytes, None]:
585    """
586    Given a serialized ASCII string of bytes data, return the original bytes.
587    The input data may either be base64- or hex-encoded.
588
589    Parameters
590    ----------
591    data: str | None
592        The string to be deserialized into bytes.
593        May be base64- or hex-encoded (prefixed with `'\\x'`).
594
595    force_hex: bool = False
596        If `True`, treat the input string as hex-encoded.
597        If `data` does not begin with the prefix `'\\x'`, set `force_hex` to `True`.
598        This will still strip the leading `'\\x'` prefix if present.
599
600    Returns
601    -------
602    The original bytes used to produce the encoded string `data`.
603    """
604    if not isinstance(data, str) and value_is_null(data):
605        return data
606
607    import binascii
608    import base64
609
610    is_hex = force_hex or data.startswith('\\x')
611
612    if is_hex:
613        if data.startswith('\\x'):
614            data = data[2:]
615        return binascii.unhexlify(data)
616
617    return base64.b64decode(data)
618
619
620def deserialize_base64(data: str) -> bytes:
621    """
622    Return the original bytestring from the given base64-encoded string.
623    """
624    import base64
625    return base64.b64decode(data)
626
627
628def encode_bytes_for_bytea(data: bytes, with_prefix: bool = True) -> str | None:
629    """
630    Return the given bytes as a hex string for PostgreSQL's `BYTEA` type.
631    """
632    import binascii
633    if not isinstance(data, bytes) and value_is_null(data):
634        return data
635    return ('\\x' if with_prefix else '') + binascii.hexlify(data).decode('utf-8')
636
637
638def serialize_datetime(dt: datetime) -> Union[str, None]:
639    """
640    Serialize a datetime object into JSON (ISO format string).
641
642    Examples
643    --------
644    >>> import json
645    >>> from datetime import datetime
646    >>> json.dumps({'a': datetime(2022, 1, 1)}, default=json_serialize_datetime)
647    '{"a": "2022-01-01T00:00:00Z"}'
648
649    """
650    if not isinstance(dt, datetime):
651        return None
652    tz_suffix = 'Z' if dt.tzinfo is None else ''
653    return dt.isoformat() + tz_suffix
654
655
656def json_serialize_value(x: Any, default_to_str: bool = True) -> str:
657    """
658    Serialize the given value to a JSON value. Accounts for datetimes, bytes, decimals, etc.
659
660    Parameters
661    ----------
662    x: Any
663        The value to serialize.
664
665    default_to_str: bool, default True
666        If `True`, return a string of `x` if x is not a designated type.
667        Otherwise return x.
668
669    Returns
670    -------
671    A serialized version of x, or x.
672    """
673    if isinstance(x, (mrsm.Pipe, mrsm.connectors.Connector)):
674        return x.meta
675
676    if hasattr(x, 'tzinfo'):
677        return serialize_datetime(x)
678
679    if isinstance(x, bytes):
680        return serialize_bytes(x)
681
682    if isinstance(x, Decimal):
683        return serialize_decimal(x)
684
685    if 'shapely' in str(type(x)):
686        return serialize_geometry(x)
687
688    if value_is_null(x):
689        return None
690
691    return str(x) if default_to_str else x
692
693
694def get_geometry_type_srid(
695    dtype: str = 'geometry',
696    default_type: str = 'geometry',
697    default_srid: int = 4326,
698) -> Union[Tuple[str, int], Tuple[str, None]]:
699    """
700    Given the specified geometry `dtype`, return a tuple in the form (type, SRID).
701
702    Parameters
703    ----------
704    dtype: Optional[str], default None
705        Optionally provide a specific `geometry` syntax (e.g. `geometry[MultiLineString, 4326]`).
706        You may specify a supported `shapely` geometry type and an SRID in the dtype modifier:
707
708        - `Point`
709        - `LineString`
710        - `LinearRing`
711        - `Polygon`
712        - `MultiPoint`
713        - `MultiLineString`
714        - `MultiPolygon`
715        - `GeometryCollection`
716
717    Returns
718    -------
719    A tuple in the form (type, SRID).
720    Defaults to `(default_type, default_srid)`.
721
722    Examples
723    --------
724    >>> from meerschaum.utils.dtypes import get_geometry_type_srid
725    >>> get_geometry_type_srid()
726    ('geometry', 4326)
727    >>> get_geometry_type_srid('geometry[]')
728    ('geometry', 4326)
729    >>> get_geometry_type_srid('geometry[Point, 0]')
730    ('Point', 0)
731    >>> get_geometry_type_srid('geometry[0, Point]')
732    ('Point', 0)
733    >>> get_geometry_type_srid('geometry[0]')
734    ('geometry', 0)
735    >>> get_geometry_type_srid('geometry[MULTILINESTRING, 4326]')
736    ('MultiLineString', 4326)
737    >>> get_geometry_type_srid('geography')
738    ('geometry', 4326)
739    >>> get_geometry_type_srid('geography[POINT]')
740    ('Point', 4376)
741    """
742    from meerschaum.utils.misc import is_int
743    ### NOTE: PostGIS syntax must also be parsed.
744    dtype = dtype.replace('(', '[').replace(')', ']')
745    bare_dtype = dtype.split('[', maxsplit=1)[0]
746    modifier = dtype.split(bare_dtype, maxsplit=1)[-1].lstrip('[').rstrip(']')
747    if not modifier:
748        return default_type, default_srid
749
750    parts = [
751        part.split('=')[-1].strip()
752        for part in modifier.split(',')
753    ]
754    parts_casted = [
755        (
756            int(part)
757            if is_int(part)
758            else part
759        )
760        for part in parts
761    ]
762
763    srid = default_srid
764    geometry_type = default_type
765
766    for part in parts_casted:
767        if isinstance(part, int):
768            srid = part
769            break
770
771    for part in parts_casted:
772        if isinstance(part, str):
773            geometry_type = part
774            break
775
776    return geometry_type, srid