Commit f38e4735 authored by Pierre Augier's avatar Pierre Augier
Browse files

typeof and str2type (public API)

parent f348fce44732
from transonic._version import __version__
from transonic.aheadoftime import Transonic, boost
from transonic.typing import Array, NDim, Type, Union
from transonic.typing import Array, NDim, Type, Union, str2type, typeof
from transonic.compiler import wait_for_all_extensions
from transonic.justintime import jit, set_compile_jit
from transonic.util import set_compile_at_import
......@@ -17,5 +17,7 @@ __all__ = [
"jit",
"set_compile_jit",
"set_compile_at_import",
"str2type",
"typeof",
"wait_for_all_extensions",
]
......@@ -10,6 +10,7 @@ from transonic.typing import (
Dict,
DictMeta,
analyze_array_type,
typeof,
)
from transonic.backends.typing import base_type_formatter
......@@ -81,3 +82,32 @@ def test_float1():
dtype, ndim = analyze_array_type("float[:]")
assert dtype == "np.float"
assert ndim == 1
def test_typeof_simple():
assert typeof(1) is int
assert typeof(1.0) is float
assert typeof(1j) is complex
assert typeof("foo") is str
def test_typeof_list():
L = typeof([[1, 2], [3, 4]])
assert isinstance(L, ListMeta)
assert L.format_as_backend_type(base_type_formatter) == "int list list"
def test_typeof_dict():
D = typeof({"a": 0, "b": 1})
assert isinstance(D, DictMeta)
assert D.format_as_backend_type(base_type_formatter) == "str: int dict"
def test_typeof_array():
A = typeof(np.ones((2, 2)))
compare_array_types(A, Array[np.float64, "2d"])
def test_typeof_np_scalar():
T = typeof(np.ones(1)[0])
assert T is np.float64
......@@ -30,6 +30,8 @@ User API
.. autofunction:: str2type
.. autofunction:: typeof
Internal API
------------
......@@ -209,13 +211,16 @@ class ArrayMeta(type):
memview = False
params_filtered = []
for param in parameters:
if isinstance(param, (Type, type)):
if isinstance(param, (Type, type, np.dtype)):
if dtype is not None:
raise ValueError(
"Array should be defined with only one variable defining "
"the types. For more than one type, use "
"for example Type(float, int)"
)
if isinstance(param, np.dtype):
param = param.type
dtype = param
if isinstance(param, NDim):
......@@ -545,3 +550,60 @@ def str2type(str_type):
dtype = eval(dtype, {"np": np})
return Array[dtype, f"{ndim}d"]
_simple_types = (int, float, complex, str)
def typeof(obj):
"""Compute the Transonic type corresponding to a Python object
Supports:
- simple Python types (int, float, complex, str)
- homogeneous list
- homogeneous dict
- numpy scalars
- numpy arrays
"""
if isinstance(obj, _simple_types):
return type(obj)
if isinstance(obj, list):
if not obj:
raise ValueError("Can't determine the full type of an empty list")
type_elem = type(obj[0])
if not all(isinstance(elem, type_elem) for elem in obj):
raise ValueError("The list {obj} is not homogeneous in type")
return List[typeof(obj[0])]
if isinstance(obj, dict):
if not obj:
raise ValueError("Can't determine the full type of an empty dict")
key = next(iter(obj.keys()))
type_key = type(key)
if not all(isinstance(key, type_key) for key in obj.keys()):
raise ValueError("The dict {obj} is not homogeneous in type")
value = next(iter(obj.values()))
type_value = type(value)
if not all(isinstance(value, type_value) for value in obj.values()):
raise ValueError("The dict {obj} is not homogeneous in type")
return Dict[typeof(key), typeof(value)]
if isinstance(obj, np.ndarray):
# TODO: deeper analysis
return Array[obj.dtype, f"{obj.ndim}d"]
if np.isscalar(obj):
return obj.dtype.type
raise NotImplementedError(
f"Not able to determine the full type of {obj} (of type {type(obj)})"
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment