LeanderCS / sqlalchemy-fake-model / #3

Committed 01 Mar 2025 11:14AM UTC coverage: 85.455% (+12.0%) from 73.427%

Build # #3

Build Type

Pull #2

coveralls-python

Committed by

LeanderCS

Commit Message

1 | Add more functions

Pull Request Pull Request #2: 1 | Add more functions

Run Details

48 of 51 new or added lines in 5 files covered. (94.12%)

13 existing lines in 1 file now uncovered.

141 of 165 relevant lines covered (85.45%)

0.85 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

82.09

/sqlalchemy_fake_model/ModelFaker.py

import json
import random
import traceback
from datetime import date, datetime
from typing import Any, Dict, List, Optional, Union

from faker import Faker
from sqlalchemy import Column, ColumnDefault, Table
from sqlalchemy.orm import ColumnProperty, Session

from .Enum.ModelColumnTypesEnum import ModelColumnTypesEnum
from .Error.InvalidAmountError import InvalidAmountError
from .Model.ModelFakerConfig import ModelFakerConfig


class ModelFaker:
    """
    The ModelFaker class is a utility class that helps in generating fake data
    for a given SQLAlchemy model. It uses the faker library to generate fake
    data based on the column types of the model. It also handles relationships
    between models and can generate data for different relationships.
    """

    def __init__(
        self,
        model: Union[Table, ColumnProperty],
        db: Optional[Session] = None,
        faker: Optional[Faker] = None,
        config: Optional[ModelFakerConfig] = None,
    ) -> None:
        """
        Initializes the ModelFaker class with the given model, database session,

        :param model: The SQLAlchemy model for which fake data needs to be generated.
        :param db: Optional SQLAlchemy session to be used for creating fake data.
        :param faker: Optional Faker instance to be used for generating fake data.
        :param config: Optional ModelFakerConfig instance to be used for configuring the ModelFaker.
        """
        self.model = model
        self.db = db or self._get_framework_session()
        self.faker = faker or Faker()
        self.config = config or ModelFakerConfig()

    @staticmethod
    def _get_framework_session() -> Optional[Session]:
        """
        Tries to get the SQLAlchemy session from available frameworks.

        :return: The SQLAlchemy session if available.
        :raises RuntimeError: If no supported framework is installed or configured
        """
        try:
            from flask import current_app

            return current_app.extensions["sqlalchemy"].db.session
        except (ImportError, KeyError):
            pass

        try:
            from tornado.web import Application

            return Application().settings["db"]
        except (ImportError, KeyError):
            pass

        try:
            from django.conf import settings
            from sqlalchemy import create_engine
            from sqlalchemy.orm import sessionmaker

            engine = create_engine(settings.DATABASES["default"]["ENGINE"])
            return sessionmaker(bind=engine)()
        except (ImportError, KeyError, AttributeError):
            pass

        raise RuntimeError(
            "No SQLAlchemy session provided and no supported framework "
            "installed or configured."
        )

    def create(self, amount: Optional[int] = 1) -> None:
        """
        Creates the specified amount of fake data entries for the model.
        It handles exceptions and rolls back the session in case of any errors.

        :param amount: The number of fake data entries to create.
        :raises InvalidAmountError: If the amount is not an integer.
        """
        if not isinstance(amount, int):
            raise InvalidAmountError(amount)

        try:
            for _ in range(amount):
                data = {}

                for column in self.__getTableColumns():
                    if self.__shouldSkipField(column):
                        continue

                    data[column.name] = self._generateFakeData(column)

                if self.__isManyToManyRelationTable():
                    self.db.execute(self.model.insert().values(**data))

                else:
                    self.db.add(self.model(**data))

            self.db.commit()

        except Exception as e:
            self.db.rollback()
            raise RuntimeError(
                f"Failed to commit: {e} {traceback.format_exc()}"
            )

    def _generateFakeData(
        self, column: Column
    ) -> Optional[Union[str, int, bool, date, datetime, None]]:
        """
        Generates fake data for a given column based on its type.
        It handles Enum, String, Integer, Boolean, DateTime, and Date column
        types.

        :param column: The SQLAlchemy column for which fake data needs to be generated.
        :return: The fake data generated for the column.
        """
        columnType = column.type

        if column.doc:
            return str(self._generateJsonData(column.doc))

        # Enum has to be the first type to check, or otherwise it
        # uses the options of the corresponding type of the enum options
        elif isinstance(columnType, ModelColumnTypesEnum.ENUM.value):
            return random.choice(columnType.enums)

        elif column.foreign_keys:
            related_attribute = list(column.foreign_keys)[0].column.name
            return getattr(
                self.__handleRelationship(column), related_attribute
            )

        elif column.primary_key:
            return self._generatePrimitive(columnType)

        elif isinstance(columnType, ModelColumnTypesEnum.STRING.value):
            maxLength = (
                columnType.length if hasattr(columnType, "length") else 255
            )
            return self.faker.text(max_nb_chars=maxLength)

        elif isinstance(columnType, ModelColumnTypesEnum.INTEGER.value):
            info = column.info
            if not info:
                return self.faker.random_int()

            min_value = column.info.get("min", 1)
            max_value = column.info.get("max", 100)
            return self.faker.random_int(min=min_value, max=max_value)

        elif isinstance(columnType, ModelColumnTypesEnum.FLOAT.value):
            precision = getattr(columnType, "precision")
            if not precision:
                return self.faker.pyfloat()

            max_value = 10 ** (precision[0] - precision[1]) - 1
            return round(
                self.faker.pyfloat(min_value=0, max_value=max_value),
                precision[1],
            )

        elif isinstance(columnType, ModelColumnTypesEnum.BOOLEAN.value):
            return self.faker.boolean()

        elif isinstance(columnType, ModelColumnTypesEnum.DATE.value):
            return self.faker.date_object()

        elif isinstance(columnType, ModelColumnTypesEnum.DATETIME.value):
            return self.faker.date_time()

        return None

    def __handleRelationship(self, column: Column) -> Optional[Table]:
        """
        Handles the relationship of a column with another model.
        It creates a fake data entry for the parent model and returns its id.
        """
        parentModel = self.__getRelatedClass(column)

        ModelFaker(parentModel, self.db).create()

        return self.db.query(parentModel).first()

    def __isManyToManyRelationTable(self) -> bool:
        """
        Checks if the model is a many-to-many relationship table.
        """
        return not hasattr(self.model, "__table__") and not hasattr(
            self.model, "__mapper__"
        )

    def __shouldSkipField(self, column: Column) -> bool:
        """
        Checks if a column is a primary key or has a default value.
        """
        return (
            (column.primary_key and self.__isFieldAutoIncrement(column))
            or self.__hasFieldDefaultValue(column)
            or self.__isFieldNullable(column)
        )

    @staticmethod
    def __isFieldAutoIncrement(column: Column) -> bool:
        """
        Checks if a column is autoincrement.
        """
        return column.autoincrement and isinstance(
            column.type, ModelColumnTypesEnum.INTEGER.value
        )

    def __hasFieldDefaultValue(self, column: Column) -> bool:
        """
        Checks if a column has a default value.
        """
        return (
            isinstance(column.default, ColumnDefault)
            and column.default.arg is not None
            and not self.config.fill_default_fields
        )

    def __isFieldNullable(self, column: Column) -> bool:
        """
        Checks if a column is nullable.
        """
        return (
            column.nullable is not None
            and column.nullable is True
            and not self.config.fill_nullable_fields
        )

    def __getTableColumns(self) -> List[Column]:
        """
        Returns the columns of the model's table.
        """
        return (
            self.model.columns
            if self.__isManyToManyRelationTable()
            else self.model.__table__.columns
        )

    def __getRelatedClass(self, column: Column) -> Table:
        """
        Returns the related class of a column if it has
        a relationship with another model.
        """
        if (
            not self.__isManyToManyRelationTable()
            and column.name in self.model.__mapper__.relationships.keys()
        ):
            return self.model.__mapper__.relationships[
                column.key
            ].mapper.class_

        fk = list(column.foreign_keys)[0]

        return fk.column.table

    def _generateJsonData(self, docstring: str) -> Dict[str, Any]:
        """
        Generates JSON data based on the provided docstring.
        """
        json_structure = json.loads(docstring)

        return self._populateJsonStructure(json_structure)

    def _populateJsonStructure(
        self, structure: Union[Dict[str, Any], List[Any]]
    ) -> Any:
        """
        Populates the JSON structure with fake data based on the defined
        schema.
        """
        if isinstance(structure, dict):
            return {
                key: self._populateJsonStructure(value)
                if isinstance(value, (dict, list))
                else self._generatePrimitive(value)
                for key, value in structure.items()
            }

        elif isinstance(structure, list):
            return [
                self._populateJsonStructure(item)
                if isinstance(item, (dict, list))
                else self._generatePrimitive(item)
                for item in structure
            ]

        return structure

    def _generatePrimitive(self, primitive_type: str) -> Any:
        """
        Generates fake data for primitive types.
        """
        if primitive_type == "boolean":
            return self.faker.boolean()
        elif primitive_type == "datetime":
            return self.faker.date_time().isoformat()
        elif primitive_type == "date":
            return self.faker.date()
        elif primitive_type == "integer":
            return self.faker.random_int()
        elif primitive_type == "string":
            return self.faker.word()
        elif primitive_type == "float":
            return self.faker.pyfloat()
        return self.faker.word()

1	import json	1✔
2	import random	1✔
3	import traceback	1✔
4	from datetime import date, datetime	1✔
5	from typing import Any, Dict, List, Optional, Union	1✔
6
7	from faker import Faker	1✔
8	from sqlalchemy import Column, ColumnDefault, Table	1✔
9	from sqlalchemy.orm import ColumnProperty, Session	1✔
10
11	from .Enum.ModelColumnTypesEnum import ModelColumnTypesEnum	1✔
12	from .Error.InvalidAmountError import InvalidAmountError	1✔
13	from .Model.ModelFakerConfig import ModelFakerConfig	1✔
14
15
16	class ModelFaker:	1✔
17	"""
18	The ModelFaker class is a utility class that helps in generating fake data
19	for a given SQLAlchemy model. It uses the faker library to generate fake
20	data based on the column types of the model. It also handles relationships
21	between models and can generate data for different relationships.
22	"""
23
24	def __init__(	1✔
25	self,
26	model: Union[Table, ColumnProperty],
27	db: Optional[Session] = None,
28	faker: Optional[Faker] = None,
29	config: Optional[ModelFakerConfig] = None,
30	) -> None:
31	"""
32	Initializes the ModelFaker class with the given model, database session,
33
34	:param model: The SQLAlchemy model for which fake data needs to be generated.
35	:param db: Optional SQLAlchemy session to be used for creating fake data.
36	:param faker: Optional Faker instance to be used for generating fake data.
37	:param config: Optional ModelFakerConfig instance to be used for configuring the ModelFaker.
38	"""
39	self.model = model	1✔
40	self.db = db or self._get_framework_session()	1✔
41	self.faker = faker or Faker()	1✔
42	self.config = config or ModelFakerConfig()	1✔
43
44	@staticmethod	1✔
45	def _get_framework_session() -> Optional[Session]:	1✔
46	"""
47	Tries to get the SQLAlchemy session from available frameworks.
48
49	:return: The SQLAlchemy session if available.
50	:raises RuntimeError: If no supported framework is installed or configured
51	"""
52	try:	1✔
53	from flask import current_app	1✔
54
55	return current_app.extensions["sqlalchemy"].db.session	1✔
56	except (ImportError, KeyError):	×
UNCOV 57	pass	×
58
UNCOV 59	try:	×
UNCOV 60	from tornado.web import Application	×
61
62	return Application().settings["db"]	×
63	except (ImportError, KeyError):	×
UNCOV 64	pass	×
65
66	try:	×
UNCOV 67	from django.conf import settings	×
68	from sqlalchemy import create_engine	×
69	from sqlalchemy.orm import sessionmaker	×
70
UNCOV 71	engine = create_engine(settings.DATABASES["default"]["ENGINE"])	×
72	return sessionmaker(bind=engine)()	×
73	except (ImportError, KeyError, AttributeError):	×
74	pass	×
75
UNCOV 76	raise RuntimeError(	×
77	"No SQLAlchemy session provided and no supported framework "
78	"installed or configured."
79	)
80
81	def create(self, amount: Optional[int] = 1) -> None:	1✔
82	"""
83	Creates the specified amount of fake data entries for the model.
84	It handles exceptions and rolls back the session in case of any errors.
85
86	:param amount: The number of fake data entries to create.
87	:raises InvalidAmountError: If the amount is not an integer.
88	"""
89	if not isinstance(amount, int):	1✔
90	raise InvalidAmountError(amount)	1✔
91
92	try:	1✔
93	for _ in range(amount):	1✔
94	data = {}	1✔
95
96	for column in self.__getTableColumns():	1✔
97	if self.__shouldSkipField(column):	1✔
98	continue	1✔
99
100	data[column.name] = self._generateFakeData(column)	1✔
101
102	if self.__isManyToManyRelationTable():	1✔
103	self.db.execute(self.model.insert().values(**data))	1✔
104
105	else:
106	self.db.add(self.model(**data))	1✔
107
108	self.db.commit()	1✔
109
110	except Exception as e:	×
UNCOV 111	self.db.rollback()	×
NEW UNCOV 112	raise RuntimeError(	×
113	f"Failed to commit: {e} {traceback.format_exc()}"
114	)
115
116	def _generateFakeData(	1✔
117	self, column: Column
118	) -> Optional[Union[str, int, bool, date, datetime, None]]:
119	"""
120	Generates fake data for a given column based on its type.
121	It handles Enum, String, Integer, Boolean, DateTime, and Date column
122	types.
123
124	:param column: The SQLAlchemy column for which fake data needs to be generated.
125	:return: The fake data generated for the column.
126	"""
127	columnType = column.type	1✔
128
129	if column.doc:	1✔
130	return str(self._generateJsonData(column.doc))	1✔
131
132	# Enum has to be the first type to check, or otherwise it
133	# uses the options of the corresponding type of the enum options
134	elif isinstance(columnType, ModelColumnTypesEnum.ENUM.value):	1✔
UNCOV 135	return random.choice(columnType.enums)	×
136
137	elif column.foreign_keys:	1✔
138	related_attribute = list(column.foreign_keys)[0].column.name	1✔
139	return getattr(	1✔
140	self.__handleRelationship(column), related_attribute
141	)
142
143	elif column.primary_key:	1✔
144	return self._generatePrimitive(columnType)	1✔
145
146	elif isinstance(columnType, ModelColumnTypesEnum.STRING.value):	1✔
147	maxLength = (	1✔
148	columnType.length if hasattr(columnType, "length") else 255
149	)
150	return self.faker.text(max_nb_chars=maxLength)	1✔
151
152	elif isinstance(columnType, ModelColumnTypesEnum.INTEGER.value):	1✔
153	info = column.info	1✔
154	if not info:	1✔
155	return self.faker.random_int()	1✔
156
157	min_value = column.info.get("min", 1)	1✔
158	max_value = column.info.get("max", 100)	1✔
159	return self.faker.random_int(min=min_value, max=max_value)	1✔
160
161	elif isinstance(columnType, ModelColumnTypesEnum.FLOAT.value):	1✔
162	precision = getattr(columnType, "precision")	1✔
163	if not precision:	1✔
164	return self.faker.pyfloat()	1✔
165
166	max_value = 10 ** (precision[0] - precision[1]) - 1	1✔
167	return round(	1✔
168	self.faker.pyfloat(min_value=0, max_value=max_value),
169	precision[1],
170	)
171
172	elif isinstance(columnType, ModelColumnTypesEnum.BOOLEAN.value):	1✔
173	return self.faker.boolean()	1✔
174
175	elif isinstance(columnType, ModelColumnTypesEnum.DATE.value):	1✔
176	return self.faker.date_object()	1✔
177
178	elif isinstance(columnType, ModelColumnTypesEnum.DATETIME.value):	1✔
179	return self.faker.date_time()	1✔
180
181	return None	×
182
183	def __handleRelationship(self, column: Column) -> Optional[Table]:	1✔
184	"""
185	Handles the relationship of a column with another model.
186	It creates a fake data entry for the parent model and returns its id.
187	"""
188	parentModel = self.__getRelatedClass(column)	1✔
189
190	ModelFaker(parentModel, self.db).create()	1✔
191
192	return self.db.query(parentModel).first()	1✔
193
194	def __isManyToManyRelationTable(self) -> bool:	1✔
195	"""
196	Checks if the model is a many-to-many relationship table.
197	"""
198	return not hasattr(self.model, "__table__") and not hasattr(	1✔
199	self.model, "__mapper__"
200	)
201
202	def __shouldSkipField(self, column: Column) -> bool:	1✔
203	"""
204	Checks if a column is a primary key or has a default value.
205	"""
206	return (	1✔
207	(column.primary_key and self.__isFieldAutoIncrement(column))
208	or self.__hasFieldDefaultValue(column)
209	or self.__isFieldNullable(column)
210	)
211
212	@staticmethod	1✔
213	def __isFieldAutoIncrement(column: Column) -> bool:	1✔
214	"""
215	Checks if a column is autoincrement.
216	"""
217	return column.autoincrement and isinstance(	1✔
218	column.type, ModelColumnTypesEnum.INTEGER.value
219	)
220
221	def __hasFieldDefaultValue(self, column: Column) -> bool:	1✔
222	"""
223	Checks if a column has a default value.
224	"""
225	return (	1✔
226	isinstance(column.default, ColumnDefault)
227	and column.default.arg is not None
228	and not self.config.fill_default_fields
229	)
230
231	def __isFieldNullable(self, column: Column) -> bool:	1✔
232	"""
233	Checks if a column is nullable.
234	"""
235	return (	1✔
236	column.nullable is not None
237	and column.nullable is True
238	and not self.config.fill_nullable_fields
239	)
240
241	def __getTableColumns(self) -> List[Column]:	1✔
242	"""
243	Returns the columns of the model's table.
244	"""
245	return (	1✔
246	self.model.columns
247	if self.__isManyToManyRelationTable()
248	else self.model.__table__.columns
249	)
250
251	def __getRelatedClass(self, column: Column) -> Table:	1✔
252	"""
253	Returns the related class of a column if it has
254	a relationship with another model.
255	"""
256	if (	1✔
257	not self.__isManyToManyRelationTable()
258	and column.name in self.model.__mapper__.relationships.keys()
259	):
NEW UNCOV 260	return self.model.__mapper__.relationships[	×
261	column.key
262	].mapper.class_
263
264	fk = list(column.foreign_keys)[0]	1✔
265
266	return fk.column.table	1✔
267
268	def _generateJsonData(self, docstring: str) -> Dict[str, Any]:	1✔
269	"""
270	Generates JSON data based on the provided docstring.
271	"""
272	json_structure = json.loads(docstring)	1✔
273
274	return self._populateJsonStructure(json_structure)	1✔
275
276	def _populateJsonStructure(	1✔
277	self, structure: Union[Dict[str, Any], List[Any]]
278	) -> Any:
279	"""
280	Populates the JSON structure with fake data based on the defined
281	schema.
282	"""
283	if isinstance(structure, dict):	1✔
284	return {	1✔
285	key: self._populateJsonStructure(value)
286	if isinstance(value, (dict, list))
287	else self._generatePrimitive(value)
288	for key, value in structure.items()
289	}
290
291	elif isinstance(structure, list):	1✔
292	return [	1✔
293	self._populateJsonStructure(item)
294	if isinstance(item, (dict, list))
295	else self._generatePrimitive(item)
296	for item in structure
297	]
298
UNCOV 299	return structure	×
300
301	def _generatePrimitive(self, primitive_type: str) -> Any:	1✔
302	"""
303	Generates fake data for primitive types.
304	"""
305	if primitive_type == "boolean":	1✔
NEW UNCOV 306	return self.faker.boolean()	×
307	elif primitive_type == "datetime":	1✔
308	return self.faker.date_time().isoformat()	1✔
309	elif primitive_type == "date":	1✔
310	return self.faker.date()	1✔
311	elif primitive_type == "integer":	1✔
312	return self.faker.random_int()	1✔
313	elif primitive_type == "string":	1✔
314	return self.faker.word()	1✔
315	elif primitive_type == "float":	1✔
316	return self.faker.pyfloat()	1✔
317	return self.faker.word()	1✔

LeanderCS / sqlalchemy-fake-model / #3

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous