Merge pull request #19 from quantumlib/toplevel

Expose units at the top level and add type info for getitem

test_perf/test_value_performance.py

@@ -12,12 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, TYPE_CHECKING, Union
 import unittest
 
 from tunits import Value
 from .perf_testing_util import a_random_compatible_unit_val, a_random_unit_val, perf_goal
 
+if TYPE_CHECKING:
+    import numpy
+
 
 @perf_goal(avg_nanos=1800, args=[a_random_compatible_unit_val] * 2)
 def test_perf_add(a: Value, b: Value) -> Value:
@@ -35,7 +38,7 @@ def test_perf_multiply(a: Value, b: Value) -> Value:
 
 
 @perf_goal(avg_nanos=600, args=[a_random_compatible_unit_val] * 2)
-def test_perf_get_item(a: Value, b: Value) -> Value:
+def test_perf_get_item(a: Value, b: Value) -> Union[float, 'numpy.typing.NDArray[Any]']:
     return a[b]
 
 

test_perf/test_value_with_dimension_performance.py

@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, TYPE_CHECKING, Union
 import unittest
 
 from tunits.units import kilometer, inch
 from tunits import ValueWithDimension, Value
 from .perf_testing_util import a_random_value_with_dimension, perf_goal
 
 
+if TYPE_CHECKING:
+    import numpy
+
+
 # Using 1500 repeats so that each of 1119 different values in tunits.units get selected at least once.
 @perf_goal(repeats=1500, avg_nanos=1800, args=[a_random_value_with_dimension] * 2)
 def test_perf_add(a: ValueWithDimension, b: ValueWithDimension) -> ValueWithDimension:
@@ -37,7 +41,9 @@ def test_perf_multiply(a: ValueWithDimension, b: ValueWithDimension) -> Value:
 
 
 @perf_goal(repeats=1500, avg_nanos=600, args=[a_random_value_with_dimension] * 2)
-def test_perf_get_item(a: ValueWithDimension, b: ValueWithDimension) -> ValueWithDimension:
+def test_perf_get_item(
+    a: ValueWithDimension, b: ValueWithDimension
+) -> Union[float, 'numpy.typing.NDArray[Any]']:
     return a[b]
 
 

tunits/__init__.py

@@ -94,3 +94,5 @@
     WaveNumberArray,
     WaveNumber,
 )
+
+from tunits.units import *