scwrypts/py/lib/scwrypts/test/random_generator.py

from csv import writer, QUOTE_NONNUMERIC
from io import StringIO
from json import dumps
from random import randint, uniform, choice
from re import sub
from string import printable
from yaml import safe_dump

from .exceptions import NoDataTypeError, BadGeneratorTypeError


SUPPORTED_DATA_TYPES = None

DEFAULT_OPTIONS = {
        'data_types': None,
        'minimum':  0,
        'maximum': 64,
        'depth': 1,
        'character_set': None,
        'bool_nullable': False,
        'str_length': None,
        'str_minimum_length': 0,
        'str_maximum_length': 32,
        'list_length': 8,
        'set_length': 8,
        'dict_length': 8,
        'dict_key_types': {int, float, chr, str},
        'csv_columns': None,
        'csv_columns_minimum': 1,
        'csv_columns_maximum': 16,
        'csv_rows': None,
        'csv_rows_minimum': 2,
        'csv_rows_maximum': 16,
        'json_initial_type': dict,
        'yaml_initial_type': dict,
        }

def generate(data_type=None, options=None):
    '''
    generate random data with the call of a function
        use data_type to generate a single value

        use options to set generation options (key = type, value = kwargs)

        use options.data_types and omit data_type to generate a random type
    '''
    if options is None:
        options = {**DEFAULT_OPTIONS}
    else:
        options = DEFAULT_OPTIONS | options

    if data_type is None and options['data_types'] is None:
        raise NoDataTypeError()

    if data_type is None and options['data_types'] is not None:
        return generate(data_type=choice(list(options['data_types'])), options=options)

    if not isinstance(data_type, str):
        data_type = data_type.__name__

    if data_type not in Generator.get_supported_data_types():
        raise BadGeneratorTypeError(data_type)

    return getattr(Generator, f'_{data_type}')(options)

#####################################################################

class Generator:

    @classmethod
    def get_supported_data_types(cls):
        global SUPPORTED_DATA_TYPES  # pylint: disable=global-statement
        if SUPPORTED_DATA_TYPES is None:
            SUPPORTED_DATA_TYPES = {
                    sub('^_', '', data_type)
                    for data_type, method in Generator.__dict__.items()
                    if isinstance(method, staticmethod)
                    }

        return SUPPORTED_DATA_TYPES



    @staticmethod
    def _bool(options):
        return choice([True, False, None]) if options['bool_nullable'] else choice([True, False])


    @staticmethod
    def _int(options):
        return randint(options['minimum'], options['maximum'])


    @staticmethod
    def _float(options):
        return uniform(options['minimum'], options['maximum'])


    @staticmethod
    def _chr(options):
        character_set = options['character_set']
        return choice(character_set) if character_set is not None else chr(randint(0,65536))


    @staticmethod
    def _str(options):
        length = options['str_length']
        if length is None:
            length = generate(int, {
                'minimum': options['str_minimum_length'],
                'maximum': options['str_maximum_length'],
                })
        return ''.join((generate(chr, options) for _ in range(length)))


    @staticmethod
    def _list(options):
        if options['depth'] <= 0:
            return []

        options['depth'] -= 1

        if options['data_types'] is None:
            options['data_types'] = {bool, int, float, chr, str}

        return [ generate(None, options) for _ in range(options['list_length']) ]


    @staticmethod
    def _set(options):
        if options['depth'] <= 0:
            return set()

        options['depth'] -= 1

        if options['data_types'] is None:
            options['data_types'] = {bool, int, float, chr, str}

        set_options = options | {'data_types': options['data_types'] - {list, dict, set}}

        return { generate(None, set_options) for _ in range(options['set_length']) }


    @staticmethod
    def _dict(options):
        if options['depth'] <= 0:
            return {}

        options['depth'] -= 1

        if options['data_types'] is None:
            options['data_types'] = {bool, int, float, chr, str, list, set, dict}

        if len(options['data_types']) == 0:
            return {}

        key_options = options | {'data_types': options['dict_key_types']}
        return {
                generate(None, key_options): generate(None, options)
                for _ in range(options['dict_length'])
                }


    @staticmethod
    def _csv(options):
        '''
        creates a StringIO object containing csv data
        '''
        if options['data_types'] is None:
            options['data_types'] = {int, float, str}

        columns = options['csv_columns']
        if columns is None:
            columns = max(1, generate(int, {
                'minimum': options['csv_columns_minimum'],
                'maximum': options['csv_columns_maximum'],
                }))

        rows = options['csv_rows']
        if rows is None:
            rows = max(1, generate(int, {
                'minimum': options['csv_rows_minimum'],
                'maximum': options['csv_rows_maximum'],
                }))

        if options['character_set'] is None:
            options['character_set'] = printable

        csv = StringIO()
        csv_writer = writer(csv, quoting=QUOTE_NONNUMERIC)

        options['list_length'] = columns

        for line in [ generate(list, options) for _ in range(rows) ]:
            csv_writer.writerow(line)

        csv.seek(0)
        return csv


    @staticmethod
    def _json(options):
        '''
        creates a str containing json data
        '''
        if options['data_types'] is None:
            options['data_types'] = {bool, int, float, str, list, dict}

        if options['character_set'] is None:
            options['character_set'] = printable

        options['dict_key_types'] = { int, float, str }

        data = generate(options['json_initial_type'], options)
        return dumps(data)


    @staticmethod
    def _yaml(options):
        '''
        creates a StringIO object containing yaml data
        '''
        if options['data_types'] is None:
            options['data_types'] = {bool, int, float, str, list, dict}

        if options['character_set'] is None:
            options['character_set'] = printable

        options['dict_key_types'] = { int, float, str }

        yaml = StringIO()
        safe_dump(generate(options['yaml_initial_type'], options), yaml, default_flow_style=False)

        yaml.seek(0)
        return yaml


#####################################################################
if __name__ == '__main__':
    print(generate('json', {'depth': 3}))