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

Support Tuple

parent 255e56195325
from transonic._version import __version__
from transonic.aheadoftime import Transonic, boost
from transonic.typing import Array, NDim, Type, Union, str2type, typeof
from transonic.typing import (
Array,
NDim,
Type,
Union,
List,
Tuple,
Dict,
Set,
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
......@@ -14,6 +25,10 @@ __all__ = [
"NDim",
"Type",
"Union",
"List",
"Dict",
"Tuple",
"Set",
"jit",
"set_compile_jit",
"set_compile_at_import",
......
......@@ -68,6 +68,9 @@ class TypeFormatterCython(TypeFormatter):
def make_list_code(self, type_elem, **kwargs):
return "list"
def make_tuple_code(self, types, **kwargs):
return "tuple"
def memoryview_type(dtype, ndim) -> str:
end = "[" + ", ".join(":" * ndim) + "]"
......
......@@ -31,5 +31,11 @@ class TypeFormatter:
def make_list_code(self, type_elem, **kwargs):
return format_type_as_backend_type(type_elem, self, **kwargs) + " list"
def make_tuple_code(self, types, **kwargs):
strings = [
format_type_as_backend_type(type_, self, **kwargs) for type_ in types
]
return f"({', '.join(strings)})"
base_type_formatter = TypeFormatter()
......@@ -84,6 +84,13 @@ def test_set():
assert S.format_as_backend_type(base_type_formatter) == "str set"
def test_tuple():
T = str2type("(int, float[:, :])")
T.get_template_parameters()
assert repr(T) == 'Tuple[int, Array[float, "2d"]]'
assert T.format_as_backend_type(base_type_formatter) == "(int, float64[:, :])"
def test_float0():
dtype, ndim = analyze_array_type("float[]")
assert dtype == "np.float"
......
......@@ -54,6 +54,7 @@ Internal API
.. autofunction:: format_type_as_backend_type
"""
import re
import numpy as np
......@@ -164,6 +165,9 @@ class NDim(TemplateVar):
def __repr__(self):
if len(self.values) == 1 and self.shift == 0:
return f'"{self.values[0]}d"'
name = self.__name__
if self.shift < 0:
......@@ -197,7 +201,11 @@ class UnionVar(TemplateVar):
_letter = "U"
class ArrayMeta(type):
class Meta(type):
pass
class ArrayMeta(Meta):
"""Metaclass for the Array class used for type hinname_calling_module"""
def __getitem__(self, parameters):
......@@ -301,7 +309,7 @@ class ArrayMeta(type):
if self.memview:
strings.append('"memview"')
return "Array[" + ", ".join(strings) + "]"
return f"Array[{', '.join(strings)}]"
def format_as_backend_type(self, backend_type_formatter, **kwargs):
......@@ -339,7 +347,7 @@ class Array(metaclass=ArrayMeta):
"""Represent a Numpy array in type hinname_calling_module"""
class UnionMeta(type):
class UnionMeta(Meta):
"""Metaclass for the Union class"""
def __getitem__(self, types):
......@@ -359,9 +367,7 @@ class UnionMeta(type):
for type_ in self.types:
if hasattr(type_, "get_template_parameters"):
template_params.extend(type_.get_template_parameters())
template_params.append(self.template_var)
return tuple(template_params)
def __repr__(self):
......@@ -389,7 +395,7 @@ class Union(metaclass=UnionMeta):
"""
class ListMeta(type):
class ListMeta(Meta):
"""Metaclass for the List class"""
def __getitem__(self, type_elem):
......@@ -421,7 +427,7 @@ class List(metaclass=ListMeta):
"""
class DictMeta(type):
class DictMeta(Meta):
"""Metaclass for the Dict class"""
def __getitem__(self, types):
......@@ -469,7 +475,7 @@ class Dict(metaclass=DictMeta):
"""
class SetMeta(type):
class SetMeta(Meta):
"""Metaclass for the Set class"""
def __getitem__(self, type_keys):
......@@ -502,6 +508,53 @@ class Set(metaclass=SetMeta):
"""
class TupleMeta(Meta):
"""Metaclass for the Tuple class"""
def __getitem__(self, types):
if not isinstance(types, tuple):
types = (types,)
trans_types = []
for type_in in types:
if isinstance(type_in, str):
type_in(str2type(type_in))
trans_types.append(type_in)
return type("TupleBis", (Tuple,), {"types": trans_types})
def get_template_parameters(self):
template_params = []
for type_ in self.types:
if hasattr(type_, "get_template_parameters"):
template_params.extend(type_.get_template_parameters())
return tuple(template_params)
def __repr__(self):
strings = []
for type_ in self.types:
if isinstance(type_, Meta):
name = repr(type_)
elif isinstance(type_, type):
name = type_.__name__
else:
name = repr(type_)
strings.append(name)
return f"Tuple[{', '.join(strings)}]"
def format_as_backend_type(self, backend_type_formatter, **kwargs):
return backend_type_formatter.make_tuple_code(self.types, **kwargs)
class Tuple(metaclass=TupleMeta):
"""Similar to typing.Tuple
>>> T = Tuple[int, Array[int, "2d"]]
"""
def format_type_as_backend_type(type_, backend_type_formatter, **kwargs):
if isinstance(type_, str):
type_ = str2type(type_)
......@@ -536,6 +589,8 @@ def analyze_array_type(str_type):
def str2type(str_type):
str_type = str_type.strip()
if " or " in str_type:
subtypes = str_type.split(" or ")
return Union[tuple(str2type(subtype) for subtype in subtypes)]
......@@ -546,6 +601,14 @@ def str2type(str_type):
# not a simple type
pass
if str_type.startswith("(") and str_type.endswith(")"):
re_comma = re.compile(r",(?![^\[]*\])(?![^\(]*\))")
return Tuple[
tuple(
str2type(word) for word in re_comma.split(str_type[1:-1]) if word
)
]
words = [word for word in str_type.split(" ") if word]
if words[-1] == "list":
......@@ -580,8 +643,8 @@ def typeof(obj):
Supports:
- simple Python types (int, float, complex, str)
- homogeneous list
- homogeneous dict
- homogeneous list, dict and set
- tuple
- numpy scalars
- numpy arrays
......@@ -589,6 +652,9 @@ def typeof(obj):
if isinstance(obj, _simple_types):
return type(obj)
if isinstance(obj, tuple):
return Tuple[tuple(typeof(elem) for elem in obj)]
if isinstance(obj, (list, dict, set)) and not obj:
raise ValueError(
f"Cannot determine the full type of an empty {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