dammy/core.py
"""
This module contains the most important dammy classes. Ideally it should be separated
in 2 modules (core, db, dataset_generator) but the dependencies between them make this
impossible without causing circular imports
"""
import json
import csv
import random
from enum import Enum
from .iterator import Iterator
from .exceptions import DatasetRequiredException, MaximumRetriesExceededException, InvalidReferenceException, EmptyKeyException
LOCALIZATION = 'default'
############################ MISC FUNCTIONS ############################
def seed(n):
"""
Sets the seed for the random number generator in order to make all results replicable.
This function uses random.seed(), so it may alter the results of other parts of your code.
:param n: The seed to set
:type n: int
"""
random.seed(n)
############################ CORE ############################
class BaseGenerator:
DAMMY_LOCALIZATION = LOCALIZATION
"""
The base class from which all generators must inherit.
"""
def __init__(self, sql_equivalent):
self._last_generated = None
self._sql_equivalent = sql_equivalent
def iterator(self, dataset=None):
"""
Get a iterator which generates values and performs a posterior treatment on them.
By default, no treatment is done and generate_raw() is called.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A Python iterator
"""
return Iterator(self.__class__, dataset, raw=False)
def iterator_raw(self, dataset=None):
"""
Get a generator which generates values without posterior treatment.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: Python generator
:raises: NotImplementedError
"""
return Iterator(self.__class__, dataset, raw=True)
def generate_raw(self, dataset=None, localization=None):
"""
Generate without posterior treatment. All generators must implement this method.
If a generator does not implement this method a NotImplementedError will be raised
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value generated by the generator
:raises: NotImplementedError
"""
raise NotImplementedError('The generate_raw() method must be overridden')
def generate(self, dataset=None, localization=None):
"""
Generate a value and perform a posterior treatment. By default, no treatment is done
and generate_raw() is called.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value generated by the generator
"""
return self.generate_raw(dataset, localization)
def _generate(self, value):
"""
Updates the last generated value of the generator
:param value: The value to set as the last generated value
"""
self._last_generated = value
return value
def __add__(self, other):
"""
Performs the addition of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.addition, self._sql_equivalent)
def __radd__(self, other):
"""
Performs the addition of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.addition, self._sql_equivalent)
def __sub__(self, other):
"""
Performs the substraction of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.substraction, self._sql_equivalent)
def __rsub__(self, other):
"""
Performs the substraction of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.substraction, self._sql_equivalent)
def __mul__(self, other):
"""
Performs the multiplication of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.multiplication, self._sql_equivalent)
def __rmul__(self, other):
"""
Performs the multiplication of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.multiplication, self._sql_equivalent)
def __mod__(self, other):
"""
Performs the modulus of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.modulus, self._sql_equivalent)
def __rmod__(self, other):
"""
Performs the modulus of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.modulus, self._sql_equivalent)
def __floordiv__(self, other):
"""
Performs the integer division of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.integer_division, self._sql_equivalent)
def __rfloordiv__(self, other):
"""
Performs the integer division of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.integer_division, self._sql_equivalent)
def __truediv__(self, other):
"""
Performs the true division of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.division, self._sql_equivalent)
def __rtruediv__(self, other):
"""
Performs the true division of 2 BaseGenerator objects
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(other, self, OperationResult.Operator.division, self._sql_equivalent)
def __lt__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is less than the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.lower_than, self._sql_equivalent)
def __le__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is less or equal than the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.lower_equal, self._sql_equivalent)
def __eq__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is equal to the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.equal, self._sql_equivalent)
def __ne__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is not equal the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.not_equal, self._sql_equivalent)
def __gt__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is greater than the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.greater_than, self._sql_equivalent)
def __ge__(self, other):
"""
Compares 2 BaseGenerator objects to one another to check wether one is greater or equal than the other
:param other: The other generator
:type other: BaseGenerator
:returns: :class:`dammy.db.OperationResult`
"""
return OperationResult(self, other, OperationResult.Operator.greater_equal, self._sql_equivalent)
def __hash__(self):
"""
Returns the hash value of the object
:returns: int
"""
return hash((self._last_generated,))
def __str__(self):
"""
Returns the string representation of the object. By default this value corresponds to the value of a
randomly generated object without a dataset. It can raise a DatasetRequiredException if a dataset is required
:returns: string
:raises: :class:`dammy.exceptions.DatasetRequiredException`
"""
return str(self.generate())
def __getattr__(self, name):
"""
Allows getting attribute values and calling methods on generators. See AttributeGetter
:param name: The name of th attribute
:type name: string
:returns: :class:`dammy.AttributeGetter`
"""
return AttributeGetter(self, name)
class EntityGenerator(BaseGenerator):
"""
The class from which all composite generators must inherit.
.. note::
All clases inheriting from this class should not have any other methods
than the ones listed below, for the sake of consistency.
"""
def __init__(self):
items = self.__class__.__dict__.items()
self.attrs = [name for name, value in items if name[:1] != '_' and name[:6] != 'DAMMY_' and not callable(value)]
self.special_attrs = [name for name, value in items if name[:6] == 'DAMMY_']
def generate_raw(self, dataset=None, localization=None):
"""
Gets all the attributes of the class and generates a new value.
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A dict where every key value pair is an attribute and its value
:raises: :class:`dammy.exceptions.DatasetRequiredException`
"""
if localization is None:
localization = self.DAMMY_LOCALIZATION
result = {}
for attr in self.attrs:
attr_obj = getattr(self, attr)
# Get references to foreign keys and generate primary keys and unique values
if isinstance(attr_obj, ForeignKey) or isinstance(attr_obj, Unique):
generated = attr_obj.generate(dataset, localization)
for k, v in generated.items():
result[k] = v
# Generate other fields
elif isinstance(attr_obj, BaseGenerator):
result[attr] = attr_obj.generate(dataset, localization)
# Generate constant values
else:
result[attr] = attr_obj
return self._generate(result)
def _get_column_names(self):
"""
Get the names of the columns for this entity
:returns: list containing the names for each column
"""
result = []
for attr in self.attrs:
attr_obj = getattr(self, attr)
# Get references to foreign keys
if isinstance(attr_obj, ForeignKey):
result.extend(attr_obj.referenced_object.fields.keys())
# Generate primary keys and unique values
elif isinstance(attr_obj, Unique):
result.extend(attr_obj.fields.keys())
else:
result.append(attr)
return result
def _get_instances(self, number):
"""
Generate the specified number of instances
:param number: The number of instances to generate
:type number: int
:returns: list containing the specified number of instances of this entity
"""
return [self.generate() for i in range(0, number)]
def to_json(self, number, save_to=None, indent=4):
"""
Get the specified amount of instances as a list of json dicts
:param number: The number of instances
:param save_to: The path where the generated json will be saved. If none, it will be returned as a string
:type number: int
:type save_to: str
:returns: str containing the generated json if save_to=None, None in other cases
"""
if save_to is None:
return json.dumps(self._get_instances(number), indent=indent)
else:
with open(save_to, 'w') as f:
json.dump(self._get_instances(number), f, indent=indent)
f.close()
return None
def to_csv(self, number, save_to):
"""
Save the specified amount of instances in a csv file
:param number: The number of instances
:param save_to: The path to the file where the instances will be saved
:type number: int
:type save_to: str
"""
with open(save_to, 'w') as f:
writer = csv.writer(f, delimiter=',')
writer.writerow(self._get_column_names())
writer.writerows([list(r.values()) for r in self._get_instances(number)])
f.close()
############################ Generator manipulation ############################
class FunctionResult(BaseGenerator):
"""
Allows the manipulation of generators by functions
"""
def __init__(self, function, obj, *args, **kwargs):
self.obj = obj
self.function = function
self.args = args
self.kwargs = kwargs
self._sql_equivalent = obj._sql_equivalent
def generate_raw(self, dataset=None, localization=None):
"""
Generate a value and call the function using the generated value as a parameter
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The result of running the generated value through the function
"""
return self._generate(self.function(self.obj.generate(dataset), *self.args, **self.kwargs))
class AttributeGetter(BaseGenerator):
"""
Allows getting attribute values from values generated by generators
"""
def __init__(self, obj, attr):
super(AttributeGetter, self).__init__(obj._sql_equivalent)
self.obj = obj
self.attr = attr
def generate_raw(self, dataset=None, localization=None):
"""
Generate a value and get the specified attribute
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value of the attribute on the generated object
"""
return self._generate(getattr(self.obj.generate(dataset, localization), self.attr))
def __call__(self, *args, **kwargs):
"""
Call the method with the given arguments
:param *args: A list of parameters
:param **kwargs: A dictionary of parameters
:returns: :class:`dammy.MethodCaller`
"""
return MethodCaller(self.obj, self.attr, args)
class MethodCaller(BaseGenerator):
"""
Allows calling methods of values generated by generators
"""
def __init__(self, obj, method, *args, **kwargs):
super(MethodCaller, self).__init__(obj._sql_equivalent)
self.obj = obj
self.method = method
self.args = args[0]
self.kwargs = kwargs
def generate_raw(self, dataset=None, localization=None):
"""
Generate a value and call the specified method on the generated value
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value returned by the called method
"""
method = getattr(self.obj.generate(dataset, localization), self.method)
if len(self.args) == 1 and len(self.args[0]) == 0:
if len(self.kwargs) == 0:
return self._generate(method())
else:
return self._generate(method(**self.kwargs))
else:
if len(self.kwargs) == 0:
return self._generate(method(*self.args))
else:
return self._generate(method(*self.args, **self.kwargs))
class OperationResult(BaseGenerator):
"""
Allows binary operations with regular and Dammy objects
and it is returned when any of such operations is performed
"""
class Operator(Enum):
"""
Enumerated type containing all the available operators
"""
addition = 0
substraction = 1
multiplication = 2
division = 3
integer_division = 4
modulus = 5
lower_than = 6
lower_equal = 7
equal = 8
not_equal = 9
greater_than = 10
greater_equal = 11
def __init__(self, a, b, op, sql):
super(OperationResult, self).__init__(sql)
self.operator = op
self.d1 = a
self.d2 = b
@staticmethod
def _get_operand_value(op, dataset=None):
"""
Get the value of the operand. If it is a generator, the value of the operand
will be a value generated by the generator. If it is not, the value will be
the input value.
:param op: The operand value or a generator generating that value
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value returned after performing the operation
:raises: DatasetRequiredException
"""
if isinstance(op, OperationResult) or isinstance(op, AttributeGetter) or isinstance(op, FunctionResult) or isinstance(op, MethodCaller):
return op.generate_raw(dataset)
elif isinstance(op, BaseGenerator):
if op._last_generated is None:
return op.generate_raw(dataset)
else:
return op._last_generated
else:
return op
def generate_raw(self, dataset=None, localization=None):
"""
Generates a value and performs the operation.
It will raise a TypeError if the operator is invalid.
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The value returned after performing the operation
:raises: TypeError
"""
d1 = OperationResult._get_operand_value(self.d1, dataset)
d2 = OperationResult._get_operand_value(self.d2, dataset)
if self.operator == OperationResult.Operator.addition:
result = d1 + d2
elif self.operator == OperationResult.Operator.substraction:
result = d1 - d2
elif self.operator == OperationResult.Operator.multiplication:
result = d1 * d2
elif self.operator == OperationResult.Operator.division:
result = d1 / d2
elif self.operator == OperationResult.Operator.integer_division:
result = d1 // d2
elif self.operator == OperationResult.Operator.modulus:
result = d1 % d2
elif self.operator == OperationResult.Operator.lower_than:
result = d1 < d2
elif self.operator == OperationResult.Operator.lower_equal:
result = d1 <= d2
elif self.operator == OperationResult.Operator.equal:
result = d1 == d2
elif self.operator == OperationResult.Operator.not_equal:
result = d1 != d2
elif self.operator == OperationResult.Operator.greater_than:
result = d1 > d2
elif self.operator == OperationResult.Operator.greater_equal:
result = d1 >= d2
else:
raise TypeError('Unknown operator {}'.format(self.operator))
return self._generate(result)
############################ Database ############################
class AutoIncrement(BaseGenerator):
"""
Represents an automatically incrementing field. By default starts by 1 and increments by 1
"""
def __init__(self, start=1, increment=1):
super(AutoIncrement, self).__init__('INTEGER')
if self._last_generated is None:
self._last_generated = start - 1
self._increment = increment
def generate_raw(self, dataset=None, localization=None):
"""
Generates and updates the next value
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved. It will be ignored.
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: The next value of the sequence
"""
return self._generate(self._last_generated + 1)
class Unique(BaseGenerator):
"""
Represents a unique field. The generator encapsulated here, will be guaranteed to generate unique values
:param u: The generator which will generate unique values
:param max_retries: The number of times it will retry to generate the value when it has already been generated
:type u: BaseGenerator
:type max_retries: int
"""
def __init__(self, max_retries=100, **kwargs):
if len(kwargs) == 0:
raise EmptyKeyException()
self.table = None
self.generated = set()
self.max_retries = max_retries
self.fields = kwargs
def __len__(self):
return len(self.fields)
def __generate_using(self, method, dataset=None, localization=None):
generated = []
for x in self.fields.values():
generate_method = getattr(x, method)
generated.append(generate_method(dataset, localization))
generated = tuple(generated)
retries = 0
while generated in self.generated and retries < self.max_retries:
generated = []
for x in self.fields.values():
generate_method = getattr(x, method)
generated.append(generate_method(dataset, localization))
generated = tuple(generated)
retries += 1
if retries < self.max_retries:
self.generated.add(generated)
return self._generate(dict(zip(self.fields.keys(), generated)))
else:
raise MaximumRetriesExceededException(
'Maximum retries exceeded for {}. Cannot generate more than {} unique values'.format(
self.fields,
len(self.generated)
)
)
def generate_raw(self, dataset=None, localization=None):
"""
Generates a unique value
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved.
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A unique value generated by the associated generator
:raises: MaximumRetriesExceededException
"""
return self.__generate_using('generate_raw', dataset, localization)
def generate(self, dataset=None, localization=None):
"""
Generates a unique value
Implementation of the generate() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved.
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A unique value generated by the associated generator
:raises: MaximumRetriesExceededException
"""
return self.__generate_using('generate', dataset, localization)
def reset(self):
"""
Reset the uniqueness of the generator.
"""
self.generated = set()
class PrimaryKey(Unique):
"""
Represents a primary key. Every field encapsulated by this class becomes a member of the primary key. A
table cannot contain more than one primary key. This class is an alias of the Unique class, with the exception
that no more than a primary key can exist on each table, but multiple unique values are supported.
:param k: The fields which will be part of the primary key
:type k: :class:`dammy.db.BaseGenerator`
In this example the primary key of A will be the field called 'primary' which is an autoincrement field::
from dammy import EntityGenerator
from dammy.db import PrimaryKey, AutoIncrement
class A(EntityGenerator):
primary = PrimaryKey(AutoIncrement())
# More attributes...
In this other example the primary key of B will be formed by the fields called 'field1' and 'field2'::
from dammy import EntityGenerator
from dammy.db import PrimaryKey, AutoIncrement
class B(EntityGenerator):
field1 = PrimaryKey(AutoIncrement())
field2 = PrimaryKey(AutoIncrement())
# More attributes...
"""
pass
class ForeignKey(BaseGenerator):
"""
Represents a foreign key. The first parameter is the class where the referenced field is
and the second a list of strings, each of them containing the name of a field forming
the primary key. If the referenced attribut is not unique or primay key, a InvalidReferenceException is raised
:param ref_table: The table where the referenced field is
:param \*args: List of the names of the fields forming the referenced key
:type ref_table: :class:`dammy.db.EntityGenerator`
:type \*args: str
:raises: :class:`dammy.exceptions.InvalidReferenceException`
"""
def __init__(self, ref_table, ref_field):
super(ForeignKey, self).__init__(None)
attr_obj = getattr(ref_table, ref_field)
if isinstance(attr_obj, Unique):
self.referenced_table = ref_table.__name__
self.referenced_field = ref_field
self.referenced_object = attr_obj
else:
raise InvalidReferenceException('Unique or PrimaryKey expected, got {}'.format(attr_obj.__class__.__name__))
def __len__(self):
"""
Gets the size of the key
:returns: The number of fields forming the key
"""
return len(self.referenced_object)
def generate_raw(self, dataset=None, localization=None):
"""
Gets the values corresponding to the key from the given dataset. If the dataset is not specified,
a DatasetRequiredException will be raised.
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved.
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A unique value generated by the associated generator
:raises: DatasetRequiredException
"""
if dataset is None:
raise DatasetRequiredException(
'Reference to a unique field or primary key ({}) given but no dataset containing {}s supplied'.format(
self.referenced_field,
self.referenced_table
)
)
else:
if isinstance(dataset, DatasetGenerator):
while dataset._counters[self.referenced_table] != 0:
dataset._generate_entity(self.referenced_table)
chosen = random.choice(dataset[self.referenced_table])
return self._generate(dict((k, v) for k, v in chosen.items() if k in self.referenced_object.fields.keys()))
############################ dataset_generator ############################
class DatasetGenerator(BaseGenerator):
@staticmethod
def _infer_type(o):
"""
Tries to infer the SQL equivalent data type.
Raises a TypeError if no equivalent type is found
:param o: The object from which the type will be inferred
:returns: A string with the equivalent SQL type
:raises: TypeError
"""
if isinstance(o, bool):
return 'BOOLEAN'
elif isinstance(o, int):
return 'INTEGER'
elif isinstance(o, float):
return 'FLOAT'
elif isinstance(o, str):
return 'VARCHAR({})'.format(len(o))
else:
raise TypeError('Type {} has not SQL equivalent'.format(type(o)))
@staticmethod
def _sql_literal(o):
"""
Tries to convert a python literal to its SQL equivalent
:param o: The literal object to ve converted
:returns: A string containing the object as a SQL literal
"""
if isinstance(o, str):
return '"{}"'.format(o)
else:
return str(o)
"""
This class generates a set of different types of entities related to each other,
and is useful to generate entire databases, where each table is linked to another
one by a foreign key.
It can also be used in NoSQL databases.
:param \*args: Tuples where the first element is the class representing
the entity to generate and the second element is the number of entities to generate
:type \*args: tuple
"""
def __init__(self, *args):
self._fixed_counters = dict((v[0].__name__, v[1]) for v in args)
self._name_class_map = dict((v[0].__name__, v[0]) for v in args)
self._args = args
self.data = {}
self._counters = None
def _generate_entity(self, c, localization):
"""
Generates a single entity of the given class
:param c: The class that will generate the entity
:type c: :class:`dammy.BaseGenerator`
"""
if c in self.data:
self.data[c].append(self._name_class_map[c]().generate(self, localization))
self._counters[c] -= 1
else:
self.data[c] = []
self._generate_entity(c, localization)
def generate_raw(self, dataset=None, localization=None):
"""
Generate a new dataset with the previously given specifications
Implementation of the generate_raw() method from BaseGenerator.
:param dataset: The dataset from which all referenced fields will be retrieved
:type dataset: :class:`dammy.db.DatasetGenerator` or dict
:returns: A dict where every key value pair is an attribute and its value
:raises: :class:`dammy.exceptions.DatasetRequiredException`
"""
self._counters = self._fixed_counters.copy()
self.data = {}
for c, n in self._args:
if self._counters[c.__name__] != 0:
for _ in range(0, n):
self._generate_entity(c.__name__, localization)
return self._generate(self)
def to_json(self, save_to=None, indent=4):
"""
Get the JSON representation of the dataset. If a path is specified, a file is created and the resulting JSON is written
to the file. If no path is given, the generated JSON will be returned.
:param save_to: The path where the JSON will be saved
:param indent: The indentation level of the resulting JSON
:type save_to: str
:type indent: int
:returns: String containing the JSON encoded dataset or none if it has been written to a file
"""
if save_to is None:
return json.dumps(self.data, indent=indent)
else:
with open(save_to, 'w') as f:
json.dump(self.data, f, indent=indent)
f.close()
return None
def to_sql(self, save_to=None, create_tables=True):
"""
Gets the dataset as SQL INSERT statements. The generated SQL is always returned and if save_to is specified,
it is saved to that location. Additional CREATE TABLE statements are added if create_tables is set to True
:param save_to: The path where the resulting SQL will be saved.
:param create_tables: If set to true, it will generate the instructions to create the tables.
:type save_to: str
:type create_tables: bool
:returns: A string with the SQL sentences required to insert all the tuples
"""
table_order = []
tables = {}
instances = {}
for name, c in self._name_class_map.items():
tables[name] = {
'columns': [],
'column_types': [],
'constraints': []
}
instances[name] = c()
for col in instances[name].attrs:
col_obj = getattr(instances[name], col)
if isinstance(col_obj, ForeignKey):
fields = col_obj.referenced_object.fields
fk_fields = [col + '_' + name for name in fields.keys()]
fk_field_types = [x._sql_equivalent for x in fields.values()]
# Add the columns
tables[name]['columns'].extend(fk_fields)
tables[name]['column_types'].extend(fk_field_types)
# Add the constraint
tables[name]['constraints'].append(
'CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {}({})'.format(
col,
', '.join(fk_fields),
col_obj.referenced_table,
', '.join(fields.keys())
)
)
# Add the table
if col_obj.referenced_table not in table_order:
table_order.append(col_obj.referenced_table)
elif isinstance(col_obj, PrimaryKey):
# Add the columns
tables[name]['columns'].extend(col_obj.fields.keys())
tables[name]['column_types'].extend([x._sql_equivalent for x in col_obj.fields.values()])
# add the constraint
tables[name]['constraints'].append(
'CONSTRAINT {} PRIMARY KEY ({})'.format(
col,
', '.join(col_obj.fields.keys())
)
)
elif isinstance(col_obj, Unique):
# Add the columns
tables[name]['columns'].extend(col_obj.fields.keys())
tables[name]['column_types'].extend([x._sql_equivalent for x in col_obj.fields.values()])
# add the constraint
tables[name]['constraints'].append(
'CONSTRAINT {} UNIQUE ({})'.format(
col,
', '.join(col_obj.fields.keys())
)
)
elif isinstance(col_obj, BaseGenerator):
tables[name]['columns'].append(col)
tables[name]['column_types'].append(col_obj._sql_equivalent)
else:
tables[name]['columns'].append(col)
tables[name]['column_types'].append(DatasetGenerator._infer_type(col))
if name not in table_order:
table_order.append(name)
lines = []
if create_tables:
for table in table_order:
lines.append(
'CREATE TABLE IF NOT EXISTS {} (\n\t{}\n);'.format(
table,
',\n\t'.join([' '.join(x) for x in zip(tables[table]['columns'], tables[table]['column_types'])] + tables[table]['constraints'])
)
)
for table in table_order:
for row in self.data[table]:
lines.append(
'INSERT INTO {} ({}) VALUES ({});'.format(
table,
', '.join(tables[table]['columns']),
', '.join([DatasetGenerator._sql_literal(x) for x in row.values()]),
)
)
sql = '\n'.join(lines)
if save_to is not None:
with open(save_to, 'w') as f:
f.write(sql)
return sql
def __len__(self):
"""
Counts the number of tables
:returns: The number of tables on the dataset
"""
return len(self.data)
def __str__(self):
"""
Get all the data as a string
:returns: All the data contained on the dataset
"""
return str(self.data)
def __getitem__(self, key):
"""
Gets the content of a specific table as a dict
:param key: The name of the table to retrieve
:type key: str
:returns: A dictionary with all the data contained on the table
"""
return self.data[key]