main.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812
  1. """ Generation core.
  2. mixer.main
  3. ~~~~~~~~~~
  4. The module implements objects generation.
  5. :copyright: 2013 by Kirill Klenov.
  6. :license: BSD, see LICENSE for more details.
  7. """
  8. from __future__ import absolute_import, unicode_literals
  9. import warnings
  10. from types import GeneratorType
  11. import logging
  12. import traceback
  13. from collections import defaultdict
  14. from contextlib import contextmanager
  15. from copy import deepcopy
  16. from functools import partial
  17. from types import FunctionType, MethodType, BuiltinFunctionType
  18. from . import mix_types as t, _compat as _
  19. from .factory import GenFactory
  20. from ._faker import faker
  21. SKIP_VALUE = object()
  22. LOGLEVEL = logging.WARN
  23. LOGGER = logging.getLogger('mixer')
  24. if not LOGGER.handlers and not LOGGER.root.handlers:
  25. LOGGER.addHandler(logging.StreamHandler())
  26. class TypeMixerMeta(type):
  27. """ Cache typemixers by scheme. """
  28. mixers = dict()
  29. def __call__(cls, cls_type, mixer=None, factory=None, fake=True):
  30. backup = cls_type
  31. try:
  32. cls_type = cls.__load_cls(cls_type)
  33. assert cls_type
  34. except (AttributeError, AssertionError, LookupError):
  35. raise ValueError('Invalid scheme: %s' % backup)
  36. key = (mixer, cls_type, fake, factory)
  37. if key not in cls.mixers:
  38. cls.mixers[key] = super(TypeMixerMeta, cls).__call__(
  39. cls_type, mixer=mixer, factory=factory, fake=fake)
  40. return cls.mixers[key]
  41. @staticmethod
  42. def __load_cls(cls_type):
  43. if isinstance(cls_type, _.string_types):
  44. mod, cls_type = cls_type.rsplit('.', 1)
  45. mod = _.import_module(mod)
  46. cls_type = getattr(mod, cls_type)
  47. return cls_type
  48. class TypeMixer(_.with_metaclass(TypeMixerMeta)):
  49. """ Generate objects by scheme. """
  50. factory = GenFactory
  51. FAKE = property(lambda s: Mixer.FAKE)
  52. MIX = property(lambda s: Mixer.MIX)
  53. RANDOM = property(lambda s: Mixer.RANDOM)
  54. SELECT = property(lambda s: Mixer.SELECT)
  55. SKIP = property(lambda s: Mixer.SKIP)
  56. def __init__(self, cls, mixer=None, factory=None, fake=True):
  57. self.middlewares = []
  58. self.__factory = factory or self.factory
  59. self.__fake = fake
  60. self.__gen_values = defaultdict(set)
  61. self.__fabrics = dict()
  62. self.__mixer = mixer
  63. self.__scheme = cls
  64. self.__fields = _.OrderedDict(self.__load_fields())
  65. def __repr__(self):
  66. return "<TypeMixer {0}>".format(self.__scheme)
  67. def blend(self, **values):
  68. """ Generate object.
  69. :param **values: Predefined fields
  70. :return value: a generated value
  71. """
  72. defaults = deepcopy(self.__fields)
  73. # Prepare relations
  74. for key, params in values.items():
  75. if '__' in key:
  76. name, value = key.split('__', 1)
  77. if name not in defaults:
  78. defaults[name] = t.Field(None, name)
  79. defaults[name].params.update({value: params})
  80. continue
  81. defaults[key] = params
  82. values = dict(
  83. value.gen_value(self, name, value)
  84. if isinstance(value, t.ServiceValue)
  85. else self.get_value(name, value)
  86. for name, value in defaults.items()
  87. )
  88. # Parse MIX and SKIP values
  89. candidates = list(
  90. (name, value & values if isinstance(value, t.Mix) else value)
  91. for name, value in values.items()
  92. if value is not SKIP_VALUE
  93. )
  94. values = list()
  95. postprocess_values = list()
  96. for name, value in candidates:
  97. if isinstance(value, t._Deffered):
  98. postprocess_values.append((name, value))
  99. else:
  100. values.append((name, value))
  101. target = self.populate_target(values)
  102. # Run registered middlewares
  103. for middleware in self.middlewares:
  104. target = middleware(target)
  105. target = self.postprocess(target, postprocess_values)
  106. LOGGER.info('Blended: %s [%s]', target, self.__scheme) # noqa
  107. return target
  108. def postprocess(self, target, postprocess_values):
  109. """ Run the code after a generation. """
  110. if self.__mixer:
  111. target = self.__mixer.postprocess(target)
  112. for name, deffered in postprocess_values:
  113. setattr(target, name, deffered.value)
  114. return target
  115. def populate_target(self, values):
  116. """ Populate a target by values. """
  117. target = self.__scheme()
  118. for name, value in values:
  119. setattr(target, name, value)
  120. return target
  121. def get_value(self, name, value):
  122. """ Prepare value for field with name.
  123. :return : (name, value) or None
  124. """
  125. if isinstance(value, GeneratorType):
  126. return self.get_value(name, next(value))
  127. if isinstance(value, (FunctionType, MethodType, BuiltinFunctionType)):
  128. return self.get_value(name, value())
  129. return name, value
  130. def gen_field(self, field):
  131. """ Generate value by field.
  132. :param field: Instance of :class:`Field`
  133. :return : None or (name, value) for later usage
  134. """
  135. default = self.get_default(field)
  136. if default is not SKIP_VALUE:
  137. return self.get_value(field.name, default)
  138. if not self.is_required(field):
  139. return field.name, SKIP_VALUE
  140. unique = self.is_unique(field)
  141. return self.gen_value(field.name, field, unique=unique)
  142. def gen_random(self, field_name, random):
  143. """ Generate a random value for field with `field_name`.
  144. :param field_name: Name of field for generation.
  145. :param random: Instance of :class:`~mixer.main.Random`.
  146. :return : None or (name, value) for later use
  147. """
  148. if not random.scheme:
  149. random = deepcopy(self.__fields.get(field_name))
  150. elif not isinstance(random.scheme, type):
  151. return self.get_value(field_name, faker.random_element(random.choices))
  152. return self.gen_value(field_name, random, fake=False)
  153. gen_select = gen_random
  154. def gen_fake(self, field_name, fake):
  155. """ Generate a fake value for field with `field_name`.
  156. :param field_name: Name of field for generation.
  157. :param fake: Instance of :class:`~mixer.main.Fake`.
  158. :return : None or (name, value) for later use
  159. """
  160. if not fake.scheme:
  161. fake = deepcopy(self.__fields.get(field_name))
  162. return self.gen_value(field_name, fake, fake=True)
  163. def gen_value(self, field_name, field, fake=None, unique=False):
  164. """ Generate values from basic types.
  165. :return : (name, value) for later use
  166. """
  167. fake = self.__fake if fake is None else fake
  168. if not field:
  169. field = t.Field(getattr(self.__scheme, field_name, None), field_name)
  170. fab = self.get_fabric(field, field_name, fake=fake)
  171. try:
  172. value = fab()
  173. except ValueError:
  174. value = None
  175. except Exception as exc:
  176. LOGGER.exception(exc)
  177. raise ValueError("Generation for %s (%s) has been stopped. Exception: %s" % (
  178. field_name, self.__scheme.__name__, exc))
  179. if unique and value is not SKIP_VALUE:
  180. counter = 0
  181. try:
  182. while value in self.__gen_values[field_name]:
  183. value = fab()
  184. counter += 1
  185. if counter > 100:
  186. raise RuntimeError("Cannot generate a unique value for %s" % field_name)
  187. self.__gen_values[field_name].add(value)
  188. except TypeError:
  189. pass
  190. return self.get_value(field_name, value)
  191. def get_fabric(self, field, field_name=None, fake=None):
  192. """ Get an objects fabric for field and cache it.
  193. :param field: Field for looking a fabric
  194. :param field_name: Name of field for generation
  195. :param fake: Generate fake data instead of random data.
  196. :return function:
  197. """
  198. if fake is None:
  199. fake = self.__fake
  200. if field.params:
  201. return self.make_fabric(field.scheme, field_name, fake, kwargs=field.params)
  202. key = (field.scheme, field_name, fake)
  203. if key not in self.__fabrics:
  204. self.__fabrics[key] = self.make_fabric(field.scheme, field_name, fake)
  205. return self.__fabrics[key]
  206. def make_fabric(self, scheme, field_name=None, fake=None, kwargs=None): # noqa
  207. """ Make a fabric for scheme.
  208. :param field_class: Class for looking a fabric
  209. :param scheme: Scheme for generation
  210. :param fake: Generate fake data instead of random data.
  211. :return function:
  212. """
  213. kwargs = {} if kwargs is None else kwargs
  214. fab = self.__factory.get_fabric(scheme, field_name, fake)
  215. if not fab:
  216. return partial(type(self)(scheme, mixer=self.__mixer, fake=self.__fake,
  217. factory=self.__factory).blend, **kwargs)
  218. if kwargs:
  219. return partial(fab, **kwargs)
  220. return fab
  221. def register(self, field_name, func, fake=None):
  222. """ Register function as fabric for the field.
  223. :param field_name: Name of field for generation
  224. :param func: Function for data generation
  225. :param fake: Generate fake data instead of random data.
  226. ::
  227. class Scheme:
  228. id = str
  229. def func():
  230. return 'ID'
  231. mixer = TypeMixer(Scheme)
  232. mixer.register('id', func)
  233. test = mixer.blend()
  234. test.id == 'id'
  235. """
  236. if fake is None:
  237. fake = self.__fake
  238. field = self.__fields.get(field_name)
  239. if not field:
  240. return False
  241. key = (field.scheme, field_name, fake)
  242. self.__fabrics[key] = func
  243. if not isinstance(func, (FunctionType, MethodType)):
  244. self.__fabrics[key] = lambda: func
  245. @staticmethod
  246. def is_unique(field):
  247. """ Return True is field's value should be a unique.
  248. :return bool:
  249. """
  250. return False
  251. @staticmethod
  252. def is_required(field):
  253. """ Return True is field's value should be defined.
  254. :return bool:
  255. """
  256. return True
  257. @staticmethod
  258. def get_default(field):
  259. """ Return a default value for the field if it exists.
  260. :return value:
  261. """
  262. return SKIP_VALUE
  263. @staticmethod
  264. def guard(*args, **kwargs):
  265. """ Look in storage.
  266. :returns: False
  267. """
  268. return False
  269. def reload(self, obj):
  270. """ Reload the object from storage. """
  271. return deepcopy(obj)
  272. def __load_fields(self):
  273. """ Return scheme's fields. """
  274. for fname in dir(self.__scheme):
  275. if fname.startswith('_'):
  276. continue
  277. prop = getattr(self.__scheme, fname)
  278. yield fname, t.Field(prop, fname)
  279. class ProxyMixer:
  280. """ A Mixer's proxy. Using for generate more than one object.
  281. ::
  282. mixer.cycle(5).blend(somemodel)
  283. """
  284. def __init__(self, mixer, count=5, guards=None):
  285. self.count = count
  286. self.mixer = mixer
  287. self.guards = guards
  288. def blend(self, scheme, **values):
  289. """ Call :meth:`Mixer.blend` a few times. And stack results to list.
  290. :returns: A list of generated objects.
  291. """
  292. result = []
  293. if self.guards:
  294. return self.mixer._guard(scheme, self.guards, **values) # noqa
  295. for _ in range(self.count):
  296. result.append(
  297. self.mixer.blend(scheme, **values)
  298. )
  299. return result
  300. def __getattr__(self, name):
  301. raise AttributeError('Use "cycle" only for "blend"')
  302. # Support depricated attributes
  303. class _MetaMixer(type):
  304. FAKE = property(lambda cls: t.Fake())
  305. MIX = property(lambda cls: t.Mix())
  306. RANDOM = property(lambda cls: t.Random())
  307. SELECT = property(lambda cls: t.Select())
  308. SKIP = property(lambda cls: SKIP_VALUE)
  309. class Mixer(_.with_metaclass(_MetaMixer)):
  310. """ This class is using for integration to an application.
  311. :param fake: (True) Generate fake data instead of random data.
  312. :param factory: (:class:`~mixer.main.GenFactory`) Fabric's factory
  313. ::
  314. class SomeScheme:
  315. score = int
  316. name = str
  317. mixer = Mixer()
  318. instance = mixer.blend(SomeScheme)
  319. print instance.name # Some like: 'Mike Douglass'
  320. mixer = Mixer(fake=False)
  321. instance = mixer.blend(SomeScheme)
  322. print instance.name # Some like: 'AKJfdjh3'
  323. """
  324. # generator's controller class
  325. type_mixer_cls = TypeMixer
  326. def __init__(self, fake=True, factory=None, loglevel=LOGLEVEL,
  327. silence=False, locale=faker.locale, **params):
  328. """Initialize the Mixer instance.
  329. :param fake: (True) Generate fake data instead of random data.
  330. :param loglevel: ('WARN') Set level for logging
  331. :param silence: (False) Don't raise any errors if creation was falsed
  332. :param factory: (:class:`~mixer.main.GenFactory`) A class for
  333. generation values for types
  334. """
  335. self.params = params
  336. self.faker = faker
  337. self.__init_params__(fake=fake, loglevel=loglevel, silence=silence, locale=locale)
  338. self.__factory = factory or self.type_mixer_cls.factory
  339. def __getattr__(self, name):
  340. if name in ['f', 'g', 'fake', 'random', 'mix', 'select']:
  341. warnings.warn('"mixer.%s" is depricated, use "mixer.%s" instead.'
  342. % (name, name.upper()), stacklevel=2)
  343. name = name.upper()
  344. return getattr(self, name)
  345. raise AttributeError("Attribute %s not found." % name)
  346. @property
  347. def SKIP(self, *args, **kwargs):
  348. """ Do not generate a field.
  349. ::
  350. # Don't generate field 'somefield'
  351. mixer.blend(SomeScheme, somefield=mixer.skip)
  352. :returns: SKIP_VALUE
  353. """
  354. return SKIP_VALUE
  355. @property
  356. def FAKE(self, *args, **kwargs):
  357. """ Force generation of fake values. See :class:`~mixer.main.Fake`.
  358. :returns: Fake object
  359. """
  360. return self.__class__.FAKE
  361. @property
  362. def RANDOM(self, *args, **kwargs):
  363. """ Force generation of random values. See :class:`~mixer.main.Random`.
  364. :returns: Random object
  365. """
  366. return self.__class__.RANDOM
  367. @property
  368. def SELECT(self, *args, **kwargs):
  369. """ Select data from a storage. See :class:`~mixer.main.Select`.
  370. :returns: Select object
  371. """
  372. return self.__class__.SELECT
  373. @property
  374. def MIX(self, *args, **kwargs):
  375. """ Point to mixed object from future. See :class:`~mixer.main.Mix`.
  376. :returns: Mix object
  377. """
  378. return self.__class__.MIX
  379. def __init_params__(self, locale=None, **params):
  380. self.params.update(params)
  381. if locale:
  382. faker.locale = locale
  383. self.params['locale'] = faker.locale
  384. LOGGER.setLevel(self.params.get('loglevel'))
  385. def __repr__(self):
  386. return "<Mixer [{0}]>".format(
  387. 'fake' if self.params.get('fake') else 'rand')
  388. def blend(self, scheme, **values):
  389. """Generate instance of `scheme`.
  390. :param scheme: Scheme class for generation or string with class path.
  391. :param values: Keyword params with predefined values
  392. :return value: A generated instance
  393. ::
  394. mixer = Mixer()
  395. mixer.blend(SomeScheme, active=True)
  396. print scheme.active # True
  397. mixer.blend('module.SomeScheme', active=True)
  398. print scheme.active # True
  399. """
  400. type_mixer = self.get_typemixer(scheme)
  401. try:
  402. return type_mixer.blend(**values)
  403. except Exception as e:
  404. if self.params.get('silence'):
  405. return None
  406. if e.args:
  407. e.args = ('Mixer (%s): %s' % (scheme, e.args[0]),) + e.args[1:]
  408. LOGGER.error(traceback.format_exc())
  409. raise
  410. def get_typemixer(self, scheme):
  411. """ Return a cached typemixer instance.
  412. :return TypeMixer:
  413. """
  414. return self.type_mixer_cls(
  415. scheme, mixer=self,
  416. fake=self.params.get('fake'), factory=self.__factory)
  417. @staticmethod
  418. def postprocess(target):
  419. """ Run the code after generation.
  420. :return target:
  421. """
  422. return target
  423. @staticmethod # noqa
  424. def sequence(*args):
  425. """ Create a sequence for predefined values.
  426. It makes a infinity loop with given function where does increment the
  427. counter on each iteration.
  428. :param args: If method get more one arguments, them make generator
  429. from arguments (loop on arguments). If that get one
  430. argument and this equal a function, method makes
  431. a generator from them. If argument is equal string it
  432. should be using as format string.
  433. By default function is equal 'lambda x: x'.
  434. :returns: A generator
  435. Mixer can uses a generators.
  436. ::
  437. gen = (name for name in ['test0', 'test1', 'test2'])
  438. for counter in range(3):
  439. mixer.blend(Scheme, name=gen)
  440. Mixer.sequence is a helper for create generators more easy.
  441. Generate values from sequence:
  442. ::
  443. for _ in range(3):
  444. mixer.blend(Scheme, name=mixer.sequence('john', 'mike'))
  445. Make a generator from function:
  446. ::
  447. for counter in range(3):
  448. mixer.blend(Scheme, name=mixer.sequence(
  449. lambda c: 'test%s' % c
  450. ))
  451. Short format is a python formating string
  452. ::
  453. for counter in range(3):
  454. mixer.blend(Scheme, name=mixer.sequence('test{0}'))
  455. """
  456. if len(args) > 1:
  457. def gen():
  458. while True:
  459. for o in args:
  460. yield o
  461. return gen()
  462. func = args and args[0] or None
  463. if isinstance(func, _.string_types):
  464. func = func.format
  465. elif func is None:
  466. func = lambda x: x
  467. def gen2():
  468. counter = 0
  469. while True:
  470. yield func(counter)
  471. counter += 1
  472. return gen2()
  473. def cycle(self, count=5):
  474. """ Generate a few objects. The syntastic sugar for cycles.
  475. :param count: List of objects or integer.
  476. :returns: ProxyMixer
  477. ::
  478. users = mixer.cycle(5).blend('somemodule.User')
  479. profiles = mixer.cycle(5).blend(
  480. 'somemodule.Profile', user=(user for user in users)
  481. apples = mixer.cycle(10).blend(
  482. Apple, title=mixer.sequence('apple_{0}')
  483. """
  484. return ProxyMixer(self, count)
  485. def middleware(self, scheme):
  486. """ Middleware decorator.
  487. You could add the middleware layers to generation process: ::
  488. from mixer.backend.django import mixer
  489. # Register middleware to model
  490. @mixer.middleware('auth.user')
  491. def encrypt_password(user):
  492. user.set_password('test')
  493. return user
  494. You can add several middlewares.
  495. Each middleware should get one argument (generated value) and return
  496. them.
  497. """
  498. type_mixer = self.type_mixer_cls(
  499. scheme, mixer=self, fake=self.params.get('fake'),
  500. factory=self.__factory)
  501. def wrapper(middleware):
  502. type_mixer.middlewares.append(middleware)
  503. return middleware
  504. return wrapper
  505. def unregister_middleware(self, scheme, middleware):
  506. """Remove middleware from scheme
  507. """
  508. type_mixer = self.type_mixer_cls(
  509. scheme, mixer=self, fake=self.params.get('fake'),
  510. factory=self.__factory)
  511. type_mixer.middlewares.remove(middleware)
  512. def register(self, scheme, **params):
  513. """ Manualy register a function as value's generator for class.field.
  514. :param scheme: Scheme for generation (class or class path)
  515. :param params: Kwargs with generator's definitions (field_name=field_generator)
  516. ::
  517. class Scheme:
  518. id = str
  519. title = str
  520. def func():
  521. return 'ID'
  522. mixer.register(
  523. Scheme,
  524. id=func,
  525. title='Always same',
  526. )
  527. test = mixer.blend(Scheme)
  528. test.id == 'ID'
  529. test.title == 'Always same'
  530. """
  531. fake = self.params.get('fake')
  532. type_mixer = self.type_mixer_cls(
  533. scheme, mixer=self, fake=fake, factory=self.__factory)
  534. for field_name, func in params.items():
  535. type_mixer.register(field_name, func, fake=fake)
  536. # Double register for RANDOM
  537. if fake:
  538. type_mixer.register(field_name, func, fake=False)
  539. @contextmanager
  540. def ctx(self, **params):
  541. """ Redifine params for current mixer as context.
  542. ::
  543. with mixer.ctx(commit=False):
  544. hole = mixer.blend(Hole)
  545. self.assertTrue(hole)
  546. self.assertFalse(Hole.objects.count())
  547. """
  548. _params = dict((k, v) for k, v in self.params.items() if k in params)
  549. _params['locale'] = self.faker.locale
  550. try:
  551. self.__init_params__(**params)
  552. yield self
  553. finally:
  554. self.__init_params__(**_params)
  555. def reload(self, *objs):
  556. """ Reload the objects from storage. """
  557. results = []
  558. for obj in objs:
  559. scheme = type(obj)
  560. tm = self.get_typemixer(scheme)
  561. results.append(tm.reload(obj))
  562. return results if len(results) > 1 else results[0]
  563. def guard(self, *args, **kwargs):
  564. """ Abstract method. In some backends used for prevent object creation.
  565. :returns: A Proxy to mixer
  566. """
  567. return ProxyMixer(self, count=1, guards=(args, kwargs))
  568. def _guard(self, scheme, guards, **values):
  569. type_mixer = self.get_typemixer(scheme)
  570. args, kwargs = guards
  571. seek = type_mixer.guard(*args, **kwargs)
  572. if seek:
  573. LOGGER.info('Finded: %s [%s]', seek, type(seek)) # noqa
  574. return seek
  575. return self.blend(scheme, **values)
  576. # Default mixer
  577. mixer = Mixer()
  578. # pylama:ignore=E1120