Commit 9a37733faaa by Pierre Augier

Quick implementation for const in Cython

parent 82529d94083
......@@ -18,6 +18,7 @@ from transonic.typing import (
str2type,
typeof,
Optional,
const,
)
......
......@@ -73,6 +73,9 @@ class TypeFormatterCython(TypeFormatter):
def make_tuple_code(self, types, **kwargs):
return "tuple"
def make_const_code(self, code):
return "const " + code
def memoryview_type(dtype, ndim, mem_layout) -> str:
ndim_F = 0
......
import numpy as np
from transonic import Array
from transonic import Array, const
from transonic.backends import backends
backend = backends["cython"]
......@@ -33,3 +33,10 @@ def test_array():
memview,
positive_indices="positive_indices",
)
def test_const():
A = Array[int, "2d"]
assert "const " + A.format_as_backend_type(type_formatter) == const(
A
).format_as_backend_type(type_formatter)
......@@ -53,5 +53,8 @@ class TypeFormatter:
]
return f"({', '.join(strings)})"
def make_const_code(self, code):
return code
base_type_formatter = TypeFormatter()
......@@ -15,6 +15,7 @@ from transonic.typing import (
str2shape,
MemLayout,
Optional,
const,
)
from transonic.backends.typing import base_type_formatter
......@@ -156,3 +157,12 @@ def test_shape():
def test_optional():
assert repr(Optional[int]) == "Union[int, None]"
def test_const():
A = str2type("int[]")
B = const(A)
assert A.format_as_backend_type(
base_type_formatter
) == B.format_as_backend_type(base_type_formatter)
assert repr(B) == 'const(Array[int, "1d"])'
......@@ -985,3 +985,22 @@ def typeof(obj):
raise NotImplementedError(
f"Not able to determine the full type of {obj} (of type {type(obj)})"
)
class ConstType(Type):
def __init__(self, type_):
self.type = type_
def format_as_backend_type(self, backend_type_formatter, **kwargs):
return backend_type_formatter.make_const_code(
format_type_as_backend_type(
self.type, backend_type_formatter, **kwargs
)
)
def __repr__(self):
return f"const({repr(self.type)})"
def const(type_):
return ConstType(type_)
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