Read about our upcoming Code of Conduct on this issue

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

Union

parent ebac4327749b
import numpy as np
import fluidpythran as fp
from fluidpythran import Type, NDim, Array
from fluidpythran import Type, NDim, Array, Union
T = Type(int, np.complex128)
N = NDim(1, 3)
......@@ -16,7 +16,7 @@ T = Type(int, np.complex128)
@fp.boost
def compute(a: A, b: A, c: T, d: A1, e: str):
def compute(a: A, b: A, c: T, d: Union[A, A1], e: str):
print(e)
tmp = a + b
return tmp
......@@ -6,7 +6,7 @@ try:
except ImportError:
pass
from .annotation import Array, NDim, Type, Shape
from .annotation import Array, NDim, Type, Shape, Union
from .aheadoftime import FluidPythran, pythran_def, boost, make_signature, include
......@@ -25,6 +25,7 @@ __all__ = [
"NDim",
"Type",
"Shape",
"Union",
"cachedjit",
"used_by_cachedjit",
"set_compile_cachedjit",
......
......@@ -99,7 +99,10 @@ class TemplateVar:
def _check_type_values(self):
if not all(isinstance(value, self._type_values) for value in self.values):
raise TypeError
raise TypeError(
f"{self.values} "
f"{[isinstance(value, self._type_values) for value in self.values]}"
)
class Type(TemplateVar):
......@@ -172,6 +175,12 @@ class Shape(TemplateVar):
return isinstance(arg, str)
class UnionVar(TemplateVar):
_type_values = type
_letter = "U"
class ArrayMeta(type):
"""Metaclass for the Array class used for type hints"""
......@@ -244,10 +253,9 @@ class ArrayMeta(type):
return ArrayBis
def get_parameters(self):
return getattr(self, "parameters", tuple())
return getattr(self, "parameters", dict())
def get_template_parameters(self):
# FIXME: possible bug "tuple has no value method"
return tuple(
param
for param in self.get_parameters().values()
......@@ -305,7 +313,71 @@ class Array(metaclass=ArrayMeta):
pass
# FIXME: Union
class UnionMeta(type):
"""Metaclass for the Array class used for type hints"""
def __getitem__(self, types):
if not isinstance(types, tuple):
types = (types,)
fp = _get_fluidpythran_calling_module()
template_var = UnionVar(*types, _fp=fp)
return type(
"UnionBis", (Union,), {"types": types, "template_var": template_var}
)
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())
template_params.append(self.template_var)
return tuple(template_params)
def __repr__(self):
strings = []
for p in self.parameters.values():
if isinstance(p, type):
string = p.__name__
else:
string = repr(p)
strings.append(string)
return "Union[" + ", ".join(strings) + "]"
def get_pythran_type(self, **kwargs):
type_ = kwargs.pop(self.template_var.__name__)
return compute_pythran_type_from_type(type_, **kwargs)
class Union(metaclass=UnionMeta):
"""Similar to typing.Union"""
pass
def compute_pythran_type_from_type(type_, **kwargs):
if hasattr(type_, "get_pythran_type"):
return type_.get_pythran_type(**kwargs)
elif hasattr(type_, "__name__"):
return type_.__name__
else:
pythran_type = str(type_)
types = pythran_type.split(" or ")
new_types = []
for _type in types:
if "][" in _type:
# C style: we try to rewrite it in Cython style
base, dims = _type.split("[", 1)
dims = ", ".join([_ or ":" for _ in dims[:-1].split("][")])
_type = base + "[" + dims + "]"
elif _type.endswith("[]"):
_type = _type[:-2] + "[:]"
new_types.append(_type)
return " or ".join(new_types)
def compute_pythran_types_from_types(types, **kwargs):
......@@ -314,26 +386,7 @@ def compute_pythran_types_from_types(types, **kwargs):
"""
pythran_types = []
for type_ in types:
if hasattr(type_, "get_pythran_type"):
pythran_type = type_.get_pythran_type(**kwargs)
elif hasattr(type_, "__name__"):
pythran_type = type_.__name__
else:
pythran_type = str(type_)
types = pythran_type.split(" or ")
new_types = []
for _type in types:
if "][" in _type:
# C style: we try to rewrite it in Cython style
base, dims = _type.split("[", 1)
dims = ", ".join([_ or ":" for _ in dims[:-1].split("][")])
_type = base + "[" + dims + "]"
elif _type.endswith("[]"):
_type = _type[:-2] + "[:]"
new_types.append(_type)
pythran_type = " or ".join(new_types)
pythran_types.append(pythran_type)
pythran_types.append(compute_pythran_type_from_type(type_, **kwargs))
if "_empty" in pythran_types:
raise ValueError(
......
......@@ -2,7 +2,7 @@ import numpy as np
# pythran import numpy as np
from fluidpythran import FluidPythran, boost, include
from fluidpythran import FluidPythran, boost, include, Array, Union
# pythran def func(int, float)
......@@ -18,8 +18,11 @@ def func_tmp(arg):
return arg ** 2
A = Union[int, Array[int, "1d"]]
@boost
def func2(a: int, b: float):
def func2(a: A, b: float):
return a - func_tmp(b)
......
......@@ -43,6 +43,7 @@ if mpi.rank == 0:
if path_classes_dir.exists():
rmtree(path_classes_dir)
if path_classes_dir1.exists():
rmtree(path_classes_dir1)
mpi.barrier()
......
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