bazarr/libs/knowit/serializer.py

import datetime
import json
import re
import typing
from datetime import timedelta
from decimal import Decimal

import babelfish
import yaml
from yaml.composer import Composer
from yaml.constructor import SafeConstructor
from yaml.parser import Parser
from yaml.reader import Reader
from yaml.resolver import Resolver as DefaultResolver
from yaml.scanner import Scanner

from knowit.units import units
from knowit.utils import round_decimal


def format_property(profile: str, o):
    """Convert properties to string."""
    if isinstance(o, timedelta):
        return format_duration(o, profile)

    if isinstance(o, babelfish.language.Language):
        return format_language(o, profile)

    if hasattr(o, 'units'):
        return format_quantity(o, profile)

    return str(o)


def get_json_encoder(context):
    """Return json encoder that handles all needed object types."""
    class StringEncoder(json.JSONEncoder):
        """String json encoder."""

        def default(self, o):
            return format_property(context['profile'], o)

    return StringEncoder


def get_yaml_dumper(context):
    """Return yaml dumper that handles all needed object types."""
    class CustomDumper(yaml.SafeDumper):
        """Custom YAML Dumper."""

        def default_representer(self, data):
            """Convert data to string."""
            if isinstance(data, int):
                return self.represent_int(data)
            return self.represent_str(str(data))

        def default_language_representer(self, data):
            """Convert language to string."""
            return self.represent_str(format_language(data, context['profile']))

        def default_quantity_representer(self, data):
            """Convert quantity to string."""
            return self.default_representer(format_quantity(data, context['profile']))

        def default_duration_representer(self, data):
            """Convert quantity to string."""
            return self.default_representer(format_duration(data, context['profile']))

    CustomDumper.add_representer(babelfish.Language, CustomDumper.default_language_representer)
    CustomDumper.add_representer(timedelta, CustomDumper.default_duration_representer)
    CustomDumper.add_representer(units.Quantity, CustomDumper.default_quantity_representer)
    CustomDumper.add_representer(Decimal, CustomDumper.default_representer)

    return CustomDumper


def get_yaml_loader(constructors=None):
    """Return a yaml loader that handles sequences as python lists."""
    constructors = constructors or {}
    yaml_implicit_resolvers = dict(DefaultResolver.yaml_implicit_resolvers)

    class Resolver(DefaultResolver):
        """Custom YAML Resolver."""

    Resolver.yaml_implicit_resolvers.clear()
    for ch, vs in yaml_implicit_resolvers.items():
        Resolver.yaml_implicit_resolvers.setdefault(ch, []).extend(
            (tag, regexp) for tag, regexp in vs
            if not tag.endswith('float')
        )
    Resolver.add_implicit_resolver(  # regex copied from yaml source
        '!decimal',
        re.compile(r'''^(?:
            [-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+][0-9]+)?
            |\.[0-9_]+(?:[eE][-+][0-9]+)?
            |[-+]?[0-9][0-9_]*(?::[0-9]?[0-9])+\.[0-9_]*
            |[-+]?\.(?:inf|Inf|INF)
            |\.(?:nan|NaN|NAN)
        )$''', re.VERBOSE),
        list('-+0123456789.')
    )

    class CustomLoader(Reader, Scanner, Parser, Composer, SafeConstructor, Resolver):
        """Custom YAML Loader."""

        def __init__(self, stream):
            Reader.__init__(self, stream)
            Scanner.__init__(self)
            Parser.__init__(self)
            Composer.__init__(self)
            SafeConstructor.__init__(self)
            Resolver.__init__(self)

    CustomLoader.add_constructor('tag:yaml.org,2002:seq', yaml.Loader.construct_python_tuple)
    for tag, constructor in constructors.items():
        CustomLoader.add_constructor(tag, constructor)

    def decimal_constructor(loader, node):
        value = loader.construct_scalar(node)
        return Decimal(value)

    CustomLoader.add_constructor('!decimal', decimal_constructor)

    return CustomLoader


def format_duration(
        duration: datetime.timedelta,
        profile='default',
) -> typing.Union[str, Decimal]:
    if profile == 'technical':
        return str(duration)

    seconds = duration.total_seconds()
    if profile == 'code':
        return round_decimal(
            Decimal((duration.days * 86400 + duration.seconds) * 10 ** 6 + duration.microseconds) / 10**6, min_digits=1
        )

    hours = int(seconds // 3600)
    seconds = seconds - (hours * 3600)
    minutes = int(seconds // 60)
    seconds = int(seconds - (minutes * 60))
    if profile == 'human':
        if hours > 0:
            return f'{hours} hours {minutes:02d} minutes {seconds:02d} seconds'
        if minutes > 0:
            return f'{minutes} minutes {seconds:02d} seconds'
        return f'{seconds} seconds'

    return f'{hours}:{minutes:02d}:{seconds:02d}'


def format_language(
        language: babelfish.language.Language,
        profile: str = 'default',
) -> str:
    if profile in ('default', 'human'):
        return str(language.name)

    return str(language)


def format_quantity(
        quantity,
        profile='default',
) -> str:
    """Human friendly format."""
    if profile == 'code':
        return quantity.magnitude

    unit = quantity.units
    if unit != 'bit':
        technical = profile == 'technical'
        if unit == 'hertz':
            return _format_quantity(quantity.magnitude, unit='Hz', binary=technical, precision=3 if technical else 1)

        root_unit = quantity.to_root_units().units
        if root_unit == 'bit':
            return _format_quantity(quantity.magnitude, binary=technical, precision=3 if technical else 2)
        if root_unit == 'bit / second':
            return _format_quantity(quantity.magnitude, unit='bps', binary=technical, precision=3 if technical else 1)

    return str(quantity)


def _format_quantity(
        num,
        unit: str = 'B',
        binary: bool = False,
        precision: int = 2,
) -> str:
    if binary:
        factor = 1024
        affix = 'i'
    else:
        factor = 1000
        affix = ''
    for prefix in ('', 'K', 'M', 'G', 'T', 'P', 'E', 'Z'):
        if abs(num) < factor:
            break
        num /= factor
    else:
        prefix = 'Y'

    return f'{num:3.{precision}f} {prefix}{affix}{unit}'


YAMLLoader = get_yaml_loader()