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 uuid
 11from datetime import timezone, datetime
 12from decimal import Decimal, Context, InvalidOperation, ROUND_HALF_UP
 13
 14import meerschaum as mrsm
 15from meerschaum.utils.typing import Dict, Union, Any, Optional
 16from meerschaum.utils.warnings import warn
 17
 18MRSM_ALIAS_DTYPES: Dict[str, str] = {
 19    'decimal': 'numeric',
 20    'Decimal': 'numeric',
 21    'number': 'numeric',
 22    'jsonl': 'json',
 23    'JSON': 'json',
 24    'binary': 'bytes',
 25    'blob': 'bytes',
 26    'varbinary': 'bytes',
 27    'bytea': 'bytes',
 28    'guid': 'uuid',
 29    'UUID': 'uuid',
 30}
 31MRSM_PD_DTYPES: Dict[Union[str, None], str] = {
 32    'json': 'object',
 33    'numeric': 'object',
 34    'uuid': 'object',
 35    'datetime': 'datetime64[ns, UTC]',
 36    'bool': 'bool[pyarrow]',
 37    'int': 'Int64',
 38    'int8': 'Int8',
 39    'int16': 'Int16',
 40    'int32': 'Int32',
 41    'int64': 'Int64',
 42    'str': 'string[python]',
 43    'bytes': 'object',
 44    None: 'object',
 45}
 46
 47
 48def to_pandas_dtype(dtype: str) -> str:
 49    """
 50    Cast a supported Meerschaum dtype to a Pandas dtype.
 51    """
 52    known_dtype = MRSM_PD_DTYPES.get(dtype, None)
 53    if known_dtype is not None:
 54        return known_dtype
 55
 56    alias_dtype = MRSM_ALIAS_DTYPES.get(dtype, None)
 57    if alias_dtype is not None:
 58        return MRSM_PD_DTYPES[alias_dtype]
 59
 60    if dtype.startswith('numeric'):
 61        return MRSM_PD_DTYPES['numeric']
 62
 63    ### NOTE: Kind of a hack, but if the first word of the given dtype is in all caps,
 64    ### treat it as a SQL db type.
 65    if dtype.split(' ')[0].isupper():
 66        from meerschaum.utils.dtypes.sql import get_pd_type_from_db_type
 67        return get_pd_type_from_db_type(dtype)
 68
 69    from meerschaum.utils.packages import attempt_import
 70    pandas = attempt_import('pandas', lazy=False)
 71
 72    try:
 73        return str(pandas.api.types.pandas_dtype(dtype))
 74    except Exception:
 75        warn(
 76            f"Invalid dtype '{dtype}', will use 'object' instead:\n"
 77            + f"{traceback.format_exc()}",
 78            stack=False,
 79        )
 80    return 'object'
 81
 82
 83def are_dtypes_equal(
 84    ldtype: Union[str, Dict[str, str]],
 85    rdtype: Union[str, Dict[str, str]],
 86) -> bool:
 87    """
 88    Determine whether two dtype strings may be considered
 89    equivalent to avoid unnecessary conversions.
 90
 91    Parameters
 92    ----------
 93    ldtype: Union[str, Dict[str, str]]
 94        The left dtype to compare.
 95        May also provide a dtypes dictionary.
 96
 97    rdtype: Union[str, Dict[str, str]]
 98        The right dtype to compare.
 99        May also provide a dtypes dictionary.
100
101    Returns
102    -------
103    A `bool` indicating whether the two dtypes are to be considered equivalent.
104    """
105    if isinstance(ldtype, dict) and isinstance(rdtype, dict):
106        lkeys = sorted([str(k) for k in ldtype.keys()])
107        rkeys = sorted([str(k) for k in rdtype.keys()])
108        for lkey, rkey in zip(lkeys, rkeys):
109            if lkey != rkey:
110                return False
111            ltype = ldtype[lkey]
112            rtype = rdtype[rkey]
113            if not are_dtypes_equal(ltype, rtype):
114                return False
115        return True
116
117    try:
118        if ldtype == rdtype:
119            return True
120    except Exception:
121        warn(f"Exception when comparing dtypes, returning False:\n{traceback.format_exc()}")
122        return False
123
124    ### Sometimes pandas dtype objects are passed.
125    ldtype = str(ldtype).split('[', maxsplit=1)[0]
126    rdtype = str(rdtype).split('[', maxsplit=1)[0]
127
128    if ldtype in MRSM_ALIAS_DTYPES:
129        ldtype = MRSM_ALIAS_DTYPES[ldtype]
130
131    if rdtype in MRSM_ALIAS_DTYPES:
132        rdtype = MRSM_ALIAS_DTYPES[rdtype]
133
134    json_dtypes = ('json', 'object')
135    if ldtype in json_dtypes and rdtype in json_dtypes:
136        return True
137
138    numeric_dtypes = ('numeric', 'object')
139    if ldtype in numeric_dtypes and rdtype in numeric_dtypes:
140        return True
141
142    uuid_dtypes = ('uuid', 'object')
143    if ldtype in uuid_dtypes and rdtype in uuid_dtypes:
144        return True
145
146    bytes_dtypes = ('bytes', 'object')
147    if ldtype in bytes_dtypes and rdtype in bytes_dtypes:
148        return True
149
150    if ldtype.lower() == rdtype.lower():
151        return True
152
153    datetime_dtypes = ('datetime', 'timestamp')
154    ldtype_found_dt_prefix = False
155    rdtype_found_dt_prefix = False
156    for dt_prefix in datetime_dtypes:
157        ldtype_found_dt_prefix = (dt_prefix in ldtype.lower()) or ldtype_found_dt_prefix
158        rdtype_found_dt_prefix = (dt_prefix in rdtype.lower()) or rdtype_found_dt_prefix
159    if ldtype_found_dt_prefix and rdtype_found_dt_prefix:
160        return True
161
162    string_dtypes = ('str', 'string', 'object')
163    if ldtype in string_dtypes and rdtype in string_dtypes:
164        return True
165
166    int_dtypes = ('int', 'int64', 'int32', 'int16', 'int8')
167    if ldtype.lower() in int_dtypes and rdtype.lower() in int_dtypes:
168        return True
169
170    float_dtypes = ('float', 'float64', 'float32', 'float16', 'float128', 'double')
171    if ldtype.lower() in float_dtypes and rdtype.lower() in float_dtypes:
172        return True
173
174    bool_dtypes = ('bool', 'boolean')
175    if ldtype in bool_dtypes and rdtype in bool_dtypes:
176        return True
177
178    return False
179
180
181def is_dtype_numeric(dtype: str) -> bool:
182    """
183    Determine whether a given `dtype` string
184    should be considered compatible with the Meerschaum dtype `numeric`.
185
186    Parameters
187    ----------
188    dtype: str
189        The pandas-like dtype string.
190
191    Returns
192    -------
193    A bool indicating the dtype is compatible with `numeric`.
194    """
195    dtype_lower = dtype.lower()
196
197    acceptable_substrings = ('numeric', 'float', 'double', 'int')
198    for substring in acceptable_substrings:
199        if substring in dtype_lower:
200            return True
201
202    return False
203
204
205def attempt_cast_to_numeric(
206    value: Any,
207    quantize: bool = False,
208    precision: Optional[int] = None,
209    scale: Optional[int] = None,
210)-> Any:
211    """
212    Given a value, attempt to coerce it into a numeric (Decimal).
213
214    Parameters
215    ----------
216    value: Any
217        The value to be cast to a Decimal.
218
219    quantize: bool, default False
220        If `True`, quantize the decimal to the specified precision and scale.
221
222    precision: Optional[int], default None
223        If `quantize` is `True`, use this precision.
224
225    scale: Optional[int], default None
226        If `quantize` is `True`, use this scale.
227
228    Returns
229    -------
230    A `Decimal` if possible, or `value`.
231    """
232    if isinstance(value, Decimal):
233        if quantize and precision and scale:
234            return quantize_decimal(value, precision, scale)
235        return value
236    try:
237        if value_is_null(value):
238            return Decimal('NaN')
239
240        dec = Decimal(str(value))
241        if not quantize or not precision or not scale:
242            return dec
243        return quantize_decimal(dec, precision, scale)
244    except Exception:
245        return value
246
247
248def attempt_cast_to_uuid(value: Any) -> Any:
249    """
250    Given a value, attempt to coerce it into a UUID (`uuid4`).
251    """
252    if isinstance(value, uuid.UUID):
253        return value
254    try:
255        return (
256            uuid.UUID(str(value))
257            if not value_is_null(value)
258            else None
259        )
260    except Exception:
261        return value
262
263
264def attempt_cast_to_bytes(value: Any) -> Any:
265    """
266    Given a value, attempt to coerce it into a bytestring.
267    """
268    if isinstance(value, bytes):
269        return value
270    try:
271        return (
272            deserialize_bytes_string(str(value))
273            if not value_is_null(value)
274            else None
275        )
276    except Exception:
277        return value
278
279
280def value_is_null(value: Any) -> bool:
281    """
282    Determine if a value is a null-like string.
283    """
284    return str(value).lower() in ('none', 'nan', 'na', 'nat', '', '<na>')
285
286
287def none_if_null(value: Any) -> Any:
288    """
289    Return `None` if a value is a null-like string.
290    """
291    return (None if value_is_null(value) else value)
292
293
294def quantize_decimal(x: Decimal, precision: int, scale: int) -> Decimal:
295    """
296    Quantize a given `Decimal` to a known scale and precision.
297
298    Parameters
299    ----------
300    x: Decimal
301        The `Decimal` to be quantized.
302
303    precision: int
304        The total number of significant digits.
305
306    scale: int
307        The number of significant digits after the decimal point.
308
309    Returns
310    -------
311    A `Decimal` quantized to the specified scale and precision.
312    """
313    precision_decimal = Decimal(('1' * (precision - scale)) + '.' + ('1' * scale))
314    try:
315        return x.quantize(precision_decimal, context=Context(prec=precision), rounding=ROUND_HALF_UP)
316    except InvalidOperation:
317        pass
318
319    raise ValueError(f"Cannot quantize value '{x}' to {precision=}, {scale=}.")
320
321
322def serialize_decimal(
323    x: Any,
324    quantize: bool = False,
325    precision: Optional[int] = None,
326    scale: Optional[int] = None,
327) -> Any:
328    """
329    Return a quantized string of an input decimal.
330
331    Parameters
332    ----------
333    x: Any
334        The potential decimal to be serialized.
335
336    quantize: bool, default False
337        If `True`, quantize the incoming Decimal to the specified scale and precision
338        before serialization.
339
340    precision: Optional[int], default None
341        The precision of the decimal to be quantized.
342
343    scale: Optional[int], default None
344        The scale of the decimal to be quantized.
345
346    Returns
347    -------
348    A string of the input decimal or the input if not a Decimal.
349    """
350    if not isinstance(x, Decimal):
351        return x
352
353    if value_is_null(x):
354        return None
355
356    if quantize and scale and precision:
357        x = quantize_decimal(x, precision, scale)
358
359    return f"{x:f}"
360
361
362def coerce_timezone(
363    dt: Any,
364    strip_utc: bool = False,
365) -> Any:
366    """
367    Given a `datetime`, pandas `Timestamp` or `Series` of `Timestamp`,
368    return a UTC timestamp (strip timezone if `strip_utc` is `True`.
369    """
370    if dt is None:
371        return None
372
373    if isinstance(dt, int):
374        return dt
375
376    if isinstance(dt, str):
377        dateutil_parser = mrsm.attempt_import('dateutil.parser')
378        dt = dateutil_parser.parse(dt)
379
380    dt_is_series = hasattr(dt, 'dtype') and hasattr(dt, '__module__')
381
382    if dt_is_series:
383        pandas = mrsm.attempt_import('pandas', lazy=False)
384
385        if (
386            pandas.api.types.is_datetime64_any_dtype(dt) and (
387                (dt.dt.tz is not None and not strip_utc)
388                or
389                (dt.dt.tz is None and strip_utc)
390            )
391        ):
392            return dt
393
394        dt_series = to_datetime(dt, coerce_utc=False)
395        if strip_utc:
396            try:
397                if dt_series.dt.tz is not None:
398                    dt_series = dt_series.dt.tz_localize(None)
399            except Exception:
400                pass
401
402        return dt_series
403
404    if dt.tzinfo is None:
405        if strip_utc:
406            return dt
407        return dt.replace(tzinfo=timezone.utc)
408
409    utc_dt = dt.astimezone(timezone.utc)
410    if strip_utc:
411        return utc_dt.replace(tzinfo=None)
412    return utc_dt
413
414
415def to_datetime(dt_val: Any, as_pydatetime: bool = False, coerce_utc: bool = True) -> Any:
416    """
417    Wrap `pd.to_datetime()` and add support for out-of-bounds values.
418    """
419    pandas, dateutil_parser = mrsm.attempt_import('pandas', 'dateutil.parser', lazy=False)
420    is_dask = 'dask' in getattr(dt_val, '__module__', '')
421    dd = mrsm.attempt_import('dask.dataframe') if is_dask else None
422    dt_is_series = hasattr(dt_val, 'dtype') and hasattr(dt_val, '__module__')
423    pd = pandas if dd is None else dd
424
425    try:
426        new_dt_val = pd.to_datetime(dt_val, utc=True, format='ISO8601')
427        if as_pydatetime:
428            return new_dt_val.to_pydatetime()
429        return new_dt_val
430    except (pd.errors.OutOfBoundsDatetime, ValueError):
431        pass
432
433    def parse(x: Any) -> Any:
434        try:
435            return dateutil_parser.parse(x)
436        except Exception:
437            return x
438
439    if dt_is_series:
440        new_series = dt_val.apply(parse)
441        if coerce_utc:
442            return coerce_timezone(new_series)
443        return new_series
444
445    new_dt_val = parse(dt_val)
446    if not coerce_utc:
447        return new_dt_val
448    return coerce_timezone(new_dt_val)
449
450
451def serialize_bytes(data: bytes) -> str:
452    """
453    Return the given bytes as a base64-encoded string.
454    """
455    import base64
456    if not isinstance(data, bytes) and value_is_null(data):
457        return data
458    return base64.b64encode(data).decode('utf-8')
459
460
461def deserialize_bytes_string(data: str | None, force_hex: bool = False) -> bytes | None:
462    """
463    Given a serialized ASCII string of bytes data, return the original bytes.
464    The input data may either be base64- or hex-encoded.
465
466    Parameters
467    ----------
468    data: str | None
469        The string to be deserialized into bytes.
470        May be base64- or hex-encoded (prefixed with `'\\x'`).
471
472    force_hex: bool = False
473        If `True`, treat the input string as hex-encoded.
474        If `data` does not begin with the prefix `'\\x'`, set `force_hex` to `True`.
475        This will still strip the leading `'\\x'` prefix if present.
476
477    Returns
478    -------
479    The original bytes used to produce the encoded string `data`.
480    """
481    if not isinstance(data, str) and value_is_null(data):
482        return data
483
484    import binascii
485    import base64
486
487    is_hex = force_hex or data.startswith('\\x')
488
489    if is_hex:
490        if data.startswith('\\x'):
491            data = data[2:]
492        return binascii.unhexlify(data)
493
494    return base64.b64decode(data)
495
496
497def deserialize_base64(data: str) -> bytes:
498    """
499    Return the original bytestring from the given base64-encoded string.
500    """
501    import base64
502    return base64.b64decode(data)
503
504
505def encode_bytes_for_bytea(data: bytes, with_prefix: bool = True) -> str | None:
506    """
507    Return the given bytes as a hex string for PostgreSQL's `BYTEA` type.
508    """
509    import binascii
510    if not isinstance(data, bytes) and value_is_null(data):
511        return data
512    return ('\\x' if with_prefix else '') + binascii.hexlify(data).decode('utf-8')
513
514
515def serialize_datetime(dt: datetime) -> Union[str, None]:
516    """
517    Serialize a datetime object into JSON (ISO format string).
518
519    Examples
520    --------
521    >>> import json
522    >>> from datetime import datetime
523    >>> json.dumps({'a': datetime(2022, 1, 1)}, default=json_serialize_datetime)
524    '{"a": "2022-01-01T00:00:00Z"}'
525
526    """
527    if not isinstance(dt, datetime):
528        return None
529    tz_suffix = 'Z' if dt.tzinfo is None else ''
530    return dt.isoformat() + tz_suffix
531
532
533def json_serialize_value(x: Any, default_to_str: bool = True) -> str:
534    """
535    Serialize the given value to a JSON value. Accounts for datetimes, bytes, decimals, etc.
536
537    Parameters
538    ----------
539    x: Any
540        The value to serialize.
541
542    default_to_str: bool, default True
543        If `True`, return a string of `x` if x is not a designated type.
544        Otherwise return x.
545
546    Returns
547    -------
548    A serialized version of x, or x.
549    """
550    if isinstance(x, (mrsm.Pipe, mrsm.connectors.Connector)):
551        return x.meta
552
553    if hasattr(x, 'tzinfo'):
554        return serialize_datetime(x)
555
556    if isinstance(x, bytes):
557        return serialize_bytes(x)
558
559    if isinstance(x, Decimal):
560        return serialize_decimal(x)
561
562    if value_is_null(x):
563        return None
564
565    return str(x) if default_to_str else x