common.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. """
  2. Small utilities for testing.
  3. """
  4. import gc
  5. import os
  6. import sys
  7. import sysconfig
  8. from joblib._multiprocessing_helpers import mp
  9. from joblib.testing import SkipTest, skipif
  10. try:
  11. import lz4
  12. except ImportError:
  13. lz4 = None
  14. # TODO straight removal since in joblib.test.common?
  15. IS_PYPY = hasattr(sys, "pypy_version_info")
  16. IS_GIL_DISABLED = (
  17. sysconfig.get_config_var("Py_GIL_DISABLED") and not sys._is_gil_enabled()
  18. )
  19. # A decorator to run tests only when numpy is available
  20. try:
  21. import numpy as np
  22. def with_numpy(func):
  23. """A decorator to skip tests requiring numpy."""
  24. return func
  25. except ImportError:
  26. def with_numpy(func):
  27. """A decorator to skip tests requiring numpy."""
  28. def my_func():
  29. raise SkipTest("Test requires numpy")
  30. return my_func
  31. np = None
  32. # TODO: Turn this back on after refactoring yield based tests in test_hashing
  33. # with_numpy = skipif(not np, reason='Test requires numpy.')
  34. # we use memory_profiler library for memory consumption checks
  35. try:
  36. from memory_profiler import memory_usage
  37. def with_memory_profiler(func):
  38. """A decorator to skip tests requiring memory_profiler."""
  39. return func
  40. def memory_used(func, *args, **kwargs):
  41. """Compute memory usage when executing func."""
  42. gc.collect()
  43. mem_use = memory_usage((func, args, kwargs), interval=0.001)
  44. return max(mem_use) - min(mem_use)
  45. except ImportError:
  46. def with_memory_profiler(func):
  47. """A decorator to skip tests requiring memory_profiler."""
  48. def dummy_func():
  49. raise SkipTest("Test requires memory_profiler.")
  50. return dummy_func
  51. memory_usage = memory_used = None
  52. with_multiprocessing = skipif(mp is None, reason="Needs multiprocessing to run.")
  53. with_dev_shm = skipif(
  54. not os.path.exists("/dev/shm"),
  55. reason="This test requires a large /dev/shm shared memory fs.",
  56. )
  57. with_lz4 = skipif(lz4 is None, reason="Needs lz4 compression to run")
  58. without_lz4 = skipif(lz4 is not None, reason="Needs lz4 not being installed to run")