util.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280
  1. # testing/util.py
  2. # Copyright (C) 2005-2017 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: http://www.opensource.org/licenses/mit-license.php
  7. from ..util import jython, pypy, defaultdict, decorator, py2k
  8. import decimal
  9. import gc
  10. import time
  11. import random
  12. import sys
  13. import types
  14. if jython:
  15. def jython_gc_collect(*args):
  16. """aggressive gc.collect for tests."""
  17. gc.collect()
  18. time.sleep(0.1)
  19. gc.collect()
  20. gc.collect()
  21. return 0
  22. # "lazy" gc, for VM's that don't GC on refcount == 0
  23. gc_collect = lazy_gc = jython_gc_collect
  24. elif pypy:
  25. def pypy_gc_collect(*args):
  26. gc.collect()
  27. gc.collect()
  28. gc_collect = lazy_gc = pypy_gc_collect
  29. else:
  30. # assume CPython - straight gc.collect, lazy_gc() is a pass
  31. gc_collect = gc.collect
  32. def lazy_gc():
  33. pass
  34. def picklers():
  35. picklers = set()
  36. if py2k:
  37. try:
  38. import cPickle
  39. picklers.add(cPickle)
  40. except ImportError:
  41. pass
  42. import pickle
  43. picklers.add(pickle)
  44. # yes, this thing needs this much testing
  45. for pickle_ in picklers:
  46. for protocol in -1, 0, 1, 2:
  47. yield pickle_.loads, lambda d: pickle_.dumps(d, protocol)
  48. def round_decimal(value, prec):
  49. if isinstance(value, float):
  50. return round(value, prec)
  51. # can also use shift() here but that is 2.6 only
  52. return (value * decimal.Decimal("1" + "0" * prec)
  53. ).to_integral(decimal.ROUND_FLOOR) / \
  54. pow(10, prec)
  55. class RandomSet(set):
  56. def __iter__(self):
  57. l = list(set.__iter__(self))
  58. random.shuffle(l)
  59. return iter(l)
  60. def pop(self):
  61. index = random.randint(0, len(self) - 1)
  62. item = list(set.__iter__(self))[index]
  63. self.remove(item)
  64. return item
  65. def union(self, other):
  66. return RandomSet(set.union(self, other))
  67. def difference(self, other):
  68. return RandomSet(set.difference(self, other))
  69. def intersection(self, other):
  70. return RandomSet(set.intersection(self, other))
  71. def copy(self):
  72. return RandomSet(self)
  73. def conforms_partial_ordering(tuples, sorted_elements):
  74. """True if the given sorting conforms to the given partial ordering."""
  75. deps = defaultdict(set)
  76. for parent, child in tuples:
  77. deps[parent].add(child)
  78. for i, node in enumerate(sorted_elements):
  79. for n in sorted_elements[i:]:
  80. if node in deps[n]:
  81. return False
  82. else:
  83. return True
  84. def all_partial_orderings(tuples, elements):
  85. edges = defaultdict(set)
  86. for parent, child in tuples:
  87. edges[child].add(parent)
  88. def _all_orderings(elements):
  89. if len(elements) == 1:
  90. yield list(elements)
  91. else:
  92. for elem in elements:
  93. subset = set(elements).difference([elem])
  94. if not subset.intersection(edges[elem]):
  95. for sub_ordering in _all_orderings(subset):
  96. yield [elem] + sub_ordering
  97. return iter(_all_orderings(elements))
  98. def function_named(fn, name):
  99. """Return a function with a given __name__.
  100. Will assign to __name__ and return the original function if possible on
  101. the Python implementation, otherwise a new function will be constructed.
  102. This function should be phased out as much as possible
  103. in favor of @decorator. Tests that "generate" many named tests
  104. should be modernized.
  105. """
  106. try:
  107. fn.__name__ = name
  108. except TypeError:
  109. fn = types.FunctionType(fn.__code__, fn.__globals__, name,
  110. fn.__defaults__, fn.__closure__)
  111. return fn
  112. def run_as_contextmanager(ctx, fn, *arg, **kw):
  113. """Run the given function under the given contextmanager,
  114. simulating the behavior of 'with' to support older
  115. Python versions.
  116. This is not necessary anymore as we have placed 2.6
  117. as minimum Python version, however some tests are still using
  118. this structure.
  119. """
  120. obj = ctx.__enter__()
  121. try:
  122. result = fn(obj, *arg, **kw)
  123. ctx.__exit__(None, None, None)
  124. return result
  125. except:
  126. exc_info = sys.exc_info()
  127. raise_ = ctx.__exit__(*exc_info)
  128. if raise_ is None:
  129. raise
  130. else:
  131. return raise_
  132. def rowset(results):
  133. """Converts the results of sql execution into a plain set of column tuples.
  134. Useful for asserting the results of an unordered query.
  135. """
  136. return set([tuple(row) for row in results])
  137. def fail(msg):
  138. assert False, msg
  139. @decorator
  140. def provide_metadata(fn, *args, **kw):
  141. """Provide bound MetaData for a single test, dropping afterwards."""
  142. from . import config
  143. from . import engines
  144. from sqlalchemy import schema
  145. metadata = schema.MetaData(config.db)
  146. self = args[0]
  147. prev_meta = getattr(self, 'metadata', None)
  148. self.metadata = metadata
  149. try:
  150. return fn(*args, **kw)
  151. finally:
  152. engines.drop_all_tables(metadata, config.db)
  153. self.metadata = prev_meta
  154. def force_drop_names(*names):
  155. """Force the given table names to be dropped after test complete,
  156. isolating for foreign key cycles
  157. """
  158. from . import config
  159. from sqlalchemy import inspect
  160. @decorator
  161. def go(fn, *args, **kw):
  162. try:
  163. return fn(*args, **kw)
  164. finally:
  165. drop_all_tables(
  166. config.db, inspect(config.db), include_names=names)
  167. return go
  168. class adict(dict):
  169. """Dict keys available as attributes. Shadows."""
  170. def __getattribute__(self, key):
  171. try:
  172. return self[key]
  173. except KeyError:
  174. return dict.__getattribute__(self, key)
  175. def __call__(self, *keys):
  176. return tuple([self[key] for key in keys])
  177. get_all = __call__
  178. def drop_all_tables(engine, inspector, schema=None, include_names=None):
  179. from sqlalchemy import Column, Table, Integer, MetaData, \
  180. ForeignKeyConstraint
  181. from sqlalchemy.schema import DropTable, DropConstraint
  182. if include_names is not None:
  183. include_names = set(include_names)
  184. with engine.connect() as conn:
  185. for tname, fkcs in reversed(
  186. inspector.get_sorted_table_and_fkc_names(schema=schema)):
  187. if tname:
  188. if include_names is not None and tname not in include_names:
  189. continue
  190. conn.execute(DropTable(
  191. Table(tname, MetaData(), schema=schema)
  192. ))
  193. elif fkcs:
  194. if not engine.dialect.supports_alter:
  195. continue
  196. for tname, fkc in fkcs:
  197. if include_names is not None and \
  198. tname not in include_names:
  199. continue
  200. tb = Table(
  201. tname, MetaData(),
  202. Column('x', Integer),
  203. Column('y', Integer),
  204. schema=schema
  205. )
  206. conn.execute(DropConstraint(
  207. ForeignKeyConstraint(
  208. [tb.c.x], [tb.c.y], name=fkc)
  209. ))
  210. def teardown_events(event_cls):
  211. @decorator
  212. def decorate(fn, *arg, **kw):
  213. try:
  214. return fn(*arg, **kw)
  215. finally:
  216. event_cls._clear()
  217. return decorate