From 3bde28cf9191cff03ff24d69af7ac20a97dce205 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= <cr52@protonmail.com>
Date: Thu, 18 Apr 2024 22:45:40 +0200
Subject: [PATCH] RFC: refactor benchmarks setups and customize pytest's
 discovery rule to enable running benchmarks through pytest

---
 benchmarks/coordinates.py          |  9 +++++++++
 benchmarks/io_ascii/core.py        |  3 +++
 benchmarks/io_ascii/fixedwidth.py  |  3 +++
 benchmarks/io_ascii/ipac.py        |  3 +++
 benchmarks/io_ascii/main.py        |  3 +++
 benchmarks/io_ascii/rdb.py         |  3 +++
 benchmarks/io_ascii/sextractor.py  |  3 +++
 benchmarks/io_ascii/table.py       |  3 +++
 benchmarks/io_fits.py              | 18 +++++++++++++-----
 benchmarks/modeling/compound.py    |  6 ++++++
 benchmarks/stats/sigma_clipping.py |  3 +++
 benchmarks/table.py                | 12 ++++++++++++
 benchmarks/units.py                | 15 +++++++++++++++
 pytest.ini                         |  9 +++++++++
 14 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100644 pytest.ini

diff --git a/benchmarks/coordinates.py b/benchmarks/coordinates.py
index ec258d0be6..b312e0321f 100644
--- a/benchmarks/coordinates.py
+++ b/benchmarks/coordinates.py
@@ -51,6 +51,9 @@ def setup(self):
         self.array_rep = CartesianRepresentation(np.ones((3, 1000)) * u.kpc)
         self.array_dif = CartesianDifferential(np.ones((3, 1000)) * u.km / u.s)
 
+    # pytest compat
+    setup_method = setup
+
     def time_with_differentials_scalar(self):
         self.scalar_rep.with_differentials(self.scalar_dif)
 
@@ -129,6 +132,9 @@ def setup(self):
             xyz_clustered2, representation_type=CartesianRepresentation
         )
 
+    # pytest compat
+    setup_method = setup
+
     def time_init_nodata(self):
         FK5()
 
@@ -188,6 +194,9 @@ def setup(self):
             lat=self.array_q_dec, lon=self.array_q_ra
         )
 
+    # pytest compat
+    setup_method = setup
+
     def time_init_scalar(self):
         SkyCoord(1, 2, unit="deg", frame="icrs")
 
diff --git a/benchmarks/io_ascii/core.py b/benchmarks/io_ascii/core.py
index d286de4f26..a69e24f56c 100644
--- a/benchmarks/io_ascii/core.py
+++ b/benchmarks/io_ascii/core.py
@@ -41,6 +41,9 @@ def setup(self):
         for col, x in izip(self.cols, lst):
             col.str_vals = [str(s) for s in x]
 
+    # pytest compat
+    setup_method = setup
+
     def time_continuation_inputter(self):
         core.ContinuationLinesInputter().process_lines(self.lines)
 
diff --git a/benchmarks/io_ascii/fixedwidth.py b/benchmarks/io_ascii/fixedwidth.py
index 5aa5730b9a..c040aa2313 100644
--- a/benchmarks/io_ascii/fixedwidth.py
+++ b/benchmarks/io_ascii/fixedwidth.py
@@ -18,6 +18,9 @@ def setup(self):
         self.splitter.cols = self.header.cols
         self.data = ascii.FixedWidthData()
 
+    # pytest compat
+    setup_method = setup
+
     def time_splitter(self):
         self.splitter(self.lines[1:])
 
diff --git a/benchmarks/io_ascii/ipac.py b/benchmarks/io_ascii/ipac.py
index c6f7d434c4..879f400c61 100644
--- a/benchmarks/io_ascii/ipac.py
+++ b/benchmarks/io_ascii/ipac.py
@@ -26,6 +26,9 @@ def setup(self):
         self.data.cols = list(self.table.columns.values())
         self.data._set_fill_values(self.data.cols)
 
+    # pytest compat
+    setup_method = setup
+
     def time_splitter(self):
         self.splitter.join(self.vals, self.widths)
 
diff --git a/benchmarks/io_ascii/main.py b/benchmarks/io_ascii/main.py
index 391a4e4a7b..b75daeafcb 100644
--- a/benchmarks/io_ascii/main.py
+++ b/benchmarks/io_ascii/main.py
@@ -43,6 +43,9 @@ def setup(self):
         if self.file_format != "sextractor":
             self.table = self.read()
 
+    # pytest compat
+    setup_method = setup
+
     def read(self):
         return ascii.read(BytesIO(self.data), format=self.file_format, guess=False)
 
diff --git a/benchmarks/io_ascii/rdb.py b/benchmarks/io_ascii/rdb.py
index f74cfc0b72..21dea3fbee 100644
--- a/benchmarks/io_ascii/rdb.py
+++ b/benchmarks/io_ascii/rdb.py
@@ -13,5 +13,8 @@ def setup(self):
         self.lines = f.read().split("\n")
         f.close()
 
+    # pytest compat
+    setup_method = setup
+
     def time_get_cols(self):
         self.header.get_cols(self.lines)
diff --git a/benchmarks/io_ascii/sextractor.py b/benchmarks/io_ascii/sextractor.py
index 7b7018e7d2..f5ddff10e0 100644
--- a/benchmarks/io_ascii/sextractor.py
+++ b/benchmarks/io_ascii/sextractor.py
@@ -23,5 +23,8 @@ def setup(self):
             self.lines.append("# {} {} Description [pixel**2]".format(i, randword()))
         self.lines.append("Non-header line")
 
+    # pytest compat
+    setup_method = setup
+
     def time_header(self):
         self.header.get_cols(self.lines)
diff --git a/benchmarks/io_ascii/table.py b/benchmarks/io_ascii/table.py
index ad2a739dc7..b293d876b6 100644
--- a/benchmarks/io_ascii/table.py
+++ b/benchmarks/io_ascii/table.py
@@ -28,6 +28,9 @@ def setup(self):
         self.outputter = core.TableOutputter()
         self.table = table.Table()
 
+    # pytest compat
+    setup_method = setup
+
     def time_table_outputter(self):
         self.outputter(self.cols, {"table": {}})
 
diff --git a/benchmarks/io_fits.py b/benchmarks/io_fits.py
index b133f2c3fa..ab7e966d8e 100644
--- a/benchmarks/io_fits.py
+++ b/benchmarks/io_fits.py
@@ -25,6 +25,9 @@ def setup(self):
             t["booleans"] = t["floats"] > 0.5
             t.write(temp, format="fits")
 
+    # pytest compat
+    setup_method = setup
+
     def time_read_nommap(self):
         try:
             Table.read(self.table_file, format="fits", memmap=False)
@@ -42,7 +45,7 @@ def time_write(self):
         t.write(table_bytes, format="fits")
 
 
-class FITSBinTableHDU:
+class FITSBinTableHDUBenchmarks:
     def time_from_columns_bytes(self):
         x = np.repeat(b"a", 2_000_000)
         array = np.array(x, dtype=[("col", "S1")])
@@ -57,7 +60,7 @@ def make_header(ncards=1000):
     return fits.Header(cards)
 
 
-class FITSHeader:
+class FITSHeaderBenchmarks:
     """
     Tests of the Header interface
     """
@@ -66,6 +69,9 @@ def setup(self):
         self.hdr = make_header()
         self.hdr_string = self.hdr.tostring()
 
+    # pytest compat
+    setup_method = setup
+
     def time_get_int(self):
         self.hdr.get("INT999")
 
@@ -85,14 +91,13 @@ def time_fromstring(self):
         fits.Header.fromstring(self.hdr_string)
 
 
-class FITSHDUList:
+class FITSHDUListBenchmarks:
     """
     Tests of the HDUList interface
     """
 
-    filename = "many_hdu.fits"
-
     def setup_cache(self):
+        self.filename = NamedTemporaryFile(delete=False).name
         hdr = make_header()
         hdul = fits.HDUList(
             [fits.PrimaryHDU(header=hdr)]
@@ -100,6 +105,9 @@ def setup_cache(self):
         )
         hdul.writeto(self.filename)
 
+    # pytest compat
+    setup_method = setup_cache
+
     def time_getheader(self):
         fits.getheader(self.filename)
 
diff --git a/benchmarks/modeling/compound.py b/benchmarks/modeling/compound.py
index 87cd837450..ffacd952e4 100644
--- a/benchmarks/modeling/compound.py
+++ b/benchmarks/modeling/compound.py
@@ -47,6 +47,9 @@ def setup(self):
             | models.RotateNative2Celestial(5.6, -72.05, 180)
         )
 
+    # pytest compat
+    setup_method = setup
+
     def time_scalar(self):
         r, d = self.model(x_no_units_scalar, x_no_units_scalar)
 
@@ -77,6 +80,9 @@ def setup(self):
             | models.RotateNative2Celestial(5.6 * u.deg, -72.05 * u.deg, 180 * u.deg)
         )
 
+    # pytest compat
+    setup_method = setup
+
     def time_scalar(self):
         r, d = self.model(x_no_units_scalar * u.pix, x_no_units_scalar * u.pix)
 
diff --git a/benchmarks/stats/sigma_clipping.py b/benchmarks/stats/sigma_clipping.py
index 3617ff78bf..3a8238c61b 100644
--- a/benchmarks/stats/sigma_clipping.py
+++ b/benchmarks/stats/sigma_clipping.py
@@ -26,6 +26,9 @@ def setup(self):
             # deviation as the stdfunc.  The default iters is 5.
             self.sigclip = SigmaClip(sigma=3)
 
+    # pytest compat
+    setup_method = setup
+
     def time_3d_array(self):
         self.sigclip(self.data[:, :1024, :1024])
 
diff --git a/benchmarks/table.py b/benchmarks/table.py
index 6a017f9afe..c8f68d34cd 100644
--- a/benchmarks/table.py
+++ b/benchmarks/table.py
@@ -42,6 +42,9 @@ def setup(self):
 
         self.bool_mask = self.table["a"] > 0.6
 
+    # pytest compat
+    setup_method = setup
+
     def time_table_slice_bool(self):
         self.table[self.bool_mask]
 
@@ -148,6 +151,9 @@ class TimeMaskedColumn:
     def setup(self):
         self.dat = np.arange(1e7)
 
+    # pytest compat
+    setup_method = setup
+
     def time_masked_column_init(self):
         MaskedColumn(self.dat)
 
@@ -156,6 +162,9 @@ class TimeTableInitWithLists:
     def setup(self):
         self.dat = list(range(100_000))
 
+    # pytest compat
+    setup_method = setup
+
     def time_init_lists(self):
         Table([self.dat, self.dat, self.dat], names=["time", "rate", "error"])
 
@@ -183,6 +192,9 @@ def setup(self):
         self.data_str_masked_1d = self.data_str_1d.copy()
         self.data_str_masked_1d[-1] = np.ma.masked
 
+    # pytest compat
+    setup_method = setup
+
     def time_init_int_1d(self):
         Table([self.data_int_1d])
 
diff --git a/benchmarks/units.py b/benchmarks/units.py
index 91e7bda9fb..0993624af3 100644
--- a/benchmarks/units.py
+++ b/benchmarks/units.py
@@ -119,6 +119,9 @@ def setup(self):
         self.out_sq = data * u.g**2
         self.out_sqrt = data * u.g**0.5
 
+    # pytest compat
+    setup_method = setup
+
     def time_quantity_square(self):
         self.data**2
 
@@ -149,6 +152,8 @@ def setup(self):
         self.out_sq = data * u.g**2
         self.out_sqrt = data * u.g**0.5
 
+    # pytest compat
+    setup_method = setup
 
 class TimeQuantityOpSmallArrayDiffUnit:
     """
@@ -163,6 +168,9 @@ def setup(self):
         # A different but dimensionally compatible unit
         self.data2 = 0.001 * data * u.kg
 
+    # pytest compat
+    setup_method = setup
+
     def time_quantity_equal(self):
         # Same as operator.eq
         self.data == self.data2
@@ -211,6 +219,8 @@ def setup(self):
         self.data = data * u.g
         self.data2 = self.data.copy()
 
+    # pytest compat
+    setup_method = setup
 
 class TimeQuantityOpLargeArrayDiffUnit(TimeQuantityOpSmallArrayDiffUnit):
     """
@@ -224,6 +234,8 @@ def setup(self):
         # A different but dimensionally compatible unit
         self.data2 = 0.001 * data * u.kg
 
+    # pytest compat
+    setup_method = setup
 
 class TimeQuantityOpLargeArraySameUnit(TimeQuantityOpSmallArrayDiffUnit):
     """
@@ -234,3 +246,6 @@ def setup(self):
         data = np.arange(1e6) + 1
         self.data = data * u.g
         self.data2 = self.data.copy()
+
+    # pytest compat
+    setup_method = setup
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000000..e6dada0222
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,9 @@
+[pytest]
+# customize test discovery to treat benchmarks as tests
+python_files = *.py
+python_classes = Time *Benchmarks
+python_functions = time_*, timeraw_*, mem_*, peakmem_*, and track_*
+addopts =
+  # ignored because pytest interprets some method arguments as missing fixtures
+  --ignore=benchmarks/cosmology.py
+  --ignore=benchmarks/convolve.py