conftest.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. import gc
  2. import numpy as np
  3. import pytest
  4. from pandas import (
  5. DataFrame,
  6. to_datetime,
  7. )
  8. @pytest.fixture(autouse=True)
  9. def mpl_cleanup():
  10. # matplotlib/testing/decorators.py#L24
  11. # 1) Resets units registry
  12. # 2) Resets rc_context
  13. # 3) Closes all figures
  14. mpl = pytest.importorskip("matplotlib")
  15. mpl_units = pytest.importorskip("matplotlib.units")
  16. plt = pytest.importorskip("matplotlib.pyplot")
  17. orig_units_registry = mpl_units.registry.copy()
  18. with mpl.rc_context():
  19. mpl.use("template")
  20. yield
  21. mpl_units.registry.clear()
  22. mpl_units.registry.update(orig_units_registry)
  23. plt.close("all")
  24. # https://matplotlib.org/stable/users/prev_whats_new/whats_new_3.6.0.html#garbage-collection-is-no-longer-run-on-figure-close # noqa: E501
  25. gc.collect(1)
  26. @pytest.fixture
  27. def hist_df():
  28. n = 50
  29. rng = np.random.default_rng(10)
  30. gender = rng.choice(["Male", "Female"], size=n)
  31. classroom = rng.choice(["A", "B", "C"], size=n)
  32. hist_df = DataFrame(
  33. {
  34. "gender": gender,
  35. "classroom": classroom,
  36. "height": rng.normal(66, 4, size=n),
  37. "weight": rng.normal(161, 32, size=n),
  38. "category": rng.integers(4, size=n),
  39. "datetime": to_datetime(
  40. rng.integers(
  41. 812419200000000000,
  42. 819331200000000000,
  43. size=n,
  44. dtype=np.int64,
  45. )
  46. ),
  47. }
  48. )
  49. return hist_df