port benchmarks to pytest-benchmark

Author: braingram

asv.conf.json

@@ -0,0 +1,70 @@
+import io
+import itertools
+import string
+
+import numpy as np
+import pytest
+
+import asdf
+
+
+SMALL_SET = string.ascii_lowercase[:3]
+LARGE_SET = string.ascii_lowercase[:26]
+
+DATA_GENS = ["no_data", "small_data", "large_data"]
+TREE_GENS = ["small_tree",  "flat_tree", "deep_tree", "large_tree"]
+
+
+def no_data(i):
+    return i
+
+
+def small_data(i):
+    return np.full((3, 3), i)
+
+
+def large_data(i):
+    return np.full((128, 128), i)
+
+
+def small_tree(data_gen):
+    return {k: data_gen(ord(k)) for k in SMALL_SET}
+
+
+def flat_tree(data_gen):
+    return {k: data_gen(ord(k)) for k in LARGE_SET}
+
+
+def deep_tree(data_gen):
+    tree = {}
+    for k in LARGE_SET:
+        tree[k] = {"value": data_gen(k)}
+        tree = tree[k]
+    return tree
+
+
+def large_tree(data_gen):
+    tree = {}
+    for k in LARGE_SET:
+        tree["value"] = {k: data_gen(ord(k)) for k in LARGE_SET}
+        tree[k] = {}
+        tree = tree[k]
+    return tree
+
+
+@pytest.fixture(params=itertools.product(DATA_GENS, TREE_GENS))
+def tree(request):
+    data_gen_name, tree_gen_name = request.param
+    # some gymnastics to work around pytest not allowing
+    # getfixturevalue to use parametrized fixtures
+    return globals()[tree_gen_name](globals()[data_gen_name])
+
+
+@pytest.fixture
+def tree_bytes(tree):
+    return asdf.dumps(tree)
+
+
+@pytest.fixture
+def asdf_file(tree):
+    return asdf.AsdfFile(tree)

benchmarks/_utils.py

@@ -0,0 +1,35 @@
+import io
+
+import asdf
+
+
+def test_init(tree, benchmark):
+    benchmark(asdf.AsdfFile, tree)
+
+
+def test_validate(asdf_file, benchmark):
+    benchmark(asdf_file.validate)
+
+
+def test_write_to(asdf_file, benchmark):
+    bs = io.BytesIO()
+    benchmark(asdf_file.write_to, bs)
+
+
+def test_open(tree_bytes, benchmark):
+    bs = io.BytesIO(tree_bytes)
+    benchmark(asdf.open, bs)
+
+
+def test_update(tree_bytes, benchmark):
+    bs = io.BytesIO(tree_bytes)
+    af = asdf.open(bs)
+    benchmark(af.update)
+
+
+def test_dump(tree, benchmark):
+    benchmark(asdf.dumps, tree)
+
+
+def test_load(tree_bytes, benchmark):
+    benchmark(asdf.loads, tree_bytes)

benchmarks/asdf.py

@@ -0,0 +1,26 @@
+import pytest
+import numpy as np
+
+import asdf
+
+
+@pytest.fixture
+def software_asdf_file():
+    return asdf.AsdfFile({"obj": asdf.tags.core.Software(name="foo", version="0.0.0")})
+
+
+@pytest.fixture
+def ndarray_asdf_file():
+    return asdf.AsdfFile({"obj": np.ndarray([1])})
+
+
+@pytest.fixture(params=["software_asdf_file", "ndarray_asdf_file"])
+def asdf_file(request):
+    return request.getfixturevalue(request.param)
+
+
+def test_validate(asdf_file, benchmark):
+    # first validate outside the benchmark to incur
+    # extension loading, schema caching and other one-time costs
+    asdf_file.validate()
+    benchmark(asdf_file.validate)

benchmarks/conftest.py

@@ -0,0 +1,5 @@
+import asdf
+
+
+def test_walk(tree, benchmark):
+    benchmark(asdf.treeutil.walk, tree, lambda x: None)

benchmarks/schema.py

@@ -0,0 +1,21 @@
+import pytest
+
+import asdf
+
+
+@pytest.fixture(scope="module")
+def ctx_asdf_file():
+    return asdf.AsdfFile()
+
+
+@pytest.fixture()
+def tagged_tree(tree, ctx_asdf_file):
+    return asdf.yamlutil.custom_tree_to_tagged_tree(tree, ctx_asdf_file)
+
+
+def test_custom_tree_to_tagged_tree(tree, ctx_asdf_file, benchmark):
+    benchmark(asdf.yamlutil.custom_tree_to_tagged_tree, tree, ctx_asdf_file)
+
+
+def test_tagged_tree_to_tagged_tree(tagged_tree, ctx_asdf_file, benchmark):
+    benchmark(asdf.yamlutil.tagged_tree_to_custom_tree, tagged_tree, ctx_asdf_file)