From f39272501525904e6a59298e367f5c72c7e9710b Mon Sep 17 00:00:00 2001
From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com>
Date: Mon, 4 May 2026 08:40:29 -0500
Subject: [PATCH] Correct Python API signature of
 State.get_axis_values_as_strings (#346)

* Correct Python API signature of State.get_axis_values_as_strings

The C++ API has default boolean argument color, but Python API
declared no arguments.

Closes #345

* Also exercise invocation of get_axis_values_as_string with keyword argument value

* Remove use of cuda.core.experimental
---
 python/cuda/bench/__init__.pyi |  4 ++--
 python/src/py_nvbench.cpp      |  6 ++++--
 python/test/run_1.py           | 26 +++++++++++++++++++++++---
 python/test/stubs.py           |  4 ++--
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/python/cuda/bench/__init__.pyi b/python/cuda/bench/__init__.pyi
index 6b0076f..ae5d4bd 100644
--- a/python/cuda/bench/__init__.pyi
+++ b/python/cuda/bench/__init__.pyi
@@ -1,4 +1,4 @@
-# Copyright 2025 NVIDIA Corporation
+# Copyright 2025-2026 NVIDIA Corporation
 #
 #  Licensed under the Apache License, Version 2.0 with the LLVM exception
 #  (the "License"); you may not use this file except in compliance with
@@ -109,7 +109,7 @@ class State:
         self, column_name: str, value: Union[SupportsInt, SupportsFloat, str]
     ) -> None: ...
     def get_axis_values(self) -> dict[str, int | float | str]: ...
-    def get_axis_values_as_string(self) -> str: ...
+    def get_axis_values_as_string(self, color: bool = ...) -> str: ...
     def get_stopping_criterion(self) -> str: ...
 
 def register(fn: Callable[[State], None]) -> Benchmark: ...
diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp
index b1c0002..c872f62 100644
--- a/python/src/py_nvbench.cpp
+++ b/python/src/py_nvbench.cpp
@@ -1,5 +1,5 @@
 /*
- *  Copyright 2025 NVIDIA Corporation
+ *  Copyright 2025-2026 NVIDIA Corporation
  *
  *  Licensed under the Apache License, Version 2.0 with the LLVM exception
  *  (the "License"); you may not use this file except in compliance with
@@ -1097,7 +1097,9 @@ Get string of space-separated name=value pairs for this configuration
 )XXXX";
   pystate_cls.def("get_axis_values_as_string",
                   &nvbench::state::get_axis_values_as_string,
-                  method_get_axis_values_as_string_doc);
+                  method_get_axis_values_as_string_doc,
+                  py::kw_only{},
+                  py::arg("color") = false);
 
   // method State.get_axis_values
   static constexpr const char *method_get_axis_values_doc = R"XXXX(
diff --git a/python/test/run_1.py b/python/test/run_1.py
index fbc6de0..61c6a5e 100644
--- a/python/test/run_1.py
+++ b/python/test/run_1.py
@@ -1,3 +1,19 @@
+# Copyright 2026 NVIDIA Corporation
+#
+#  Licensed under the Apache License, Version 2.0 with the LLVM exception
+#  (the "License"); you may not use this file except in compliance with
+#  the License.
+#
+#  You may obtain a copy of the License at
+#
+#      http://llvm.org/foundation/relicensing/LICENSE.txt
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 import sys
 
 import cuda.bench as bench
@@ -27,6 +43,9 @@ def add_two(state: bench.State):
     a = cuda.to_device(np.random.random(N))
     c = cuda.device_array_like(a)
 
+    colorized_av = state.get_axis_values_as_string(color=True)
+    assert isinstance(colorized_av, str)
+
     assert "elements" in state.get_axis_values()
     assert "elements=" in state.get_axis_values_as_string()
 
@@ -100,9 +119,10 @@ def add_three(state: bench.State):
 
 def register_benchmarks():
     (
-        bench.register(add_two).add_int64_axis(
-            "elements", [2**pow2 - 1 for pow2 in range(20, 23)]
-        )
+        bench.register(add_two)
+        .add_int64_axis("elements", [2**pow2 - 1 for pow2 in range(20, 23)])
+        .set_stopping_criterion("entropy")
+        .set_criterion_param_float64("min-r2", 0.85)
     )
     (
         bench.register(add_float)
diff --git a/python/test/stubs.py b/python/test/stubs.py
index 0d09a58..b388607 100644
--- a/python/test/stubs.py
+++ b/python/test/stubs.py
@@ -1,4 +1,4 @@
-# Copyright 2025 NVIDIA Corporation
+# Copyright 2025-2026 NVIDIA Corporation
 #
 #  Licensed under the Apache License, Version 2.0 with the LLVM exception
 #  (the "License"); you may not use this file except in compliance with
@@ -20,7 +20,7 @@ from typing import Dict, Optional, Tuple
 
 import cuda.bench as bench
 import cuda.cccl.headers as headers
-import cuda.core.experimental as core
+import cuda.core as core
 
 
 def as_core_Stream(cs: bench.CudaStream) -> core.Stream: