draft: Generic class type alias

[paugier]

For my loop fusion demo, I see that we need a syntax compatible with generic struct for

@blue.generic
def LazySin(DTYPE, NDIM, ARG_TYPE):
    A = ndarray[DTYPE, NDIM]

    @struct
    class Self:

I propose to support

@struct
class LazySin[DTYPE, NDIM, ARG_TYPE]:
    type A = ndarray[DTYPE, NDIM]

(just syntactic sugar for the former). This PR is an implementation of that. Good shape but not yet ready for review.

Done with Claude under my supervision.

[JeffersGlass]

Just wanting to understand this a little bit - does this unblock new capabilities within generics, or is it a syntactic thing to write clearer code?

Having types as an alias only work within generic functions is just a bit confusing for me...

[antocuni]

@paugier I don't understand how type A = ndarray[DTYPE, NDIM] would be any different than A = ndarray[DTYPE, NDIM]?

E.g., the following seems to work fine:

@struct
class Foo[T]:
    TUP = tuple[T, T]

    def bar(self) -> TUP:
        return 1, 2


def main() -> None:
    f = Foo[int]()
    print(f.bar())

[paugier]

This PR is just syntactic sugar (just avoiding a lot of repetitions of potentially relatively long concrete types, typically ndarray[DTYPE, NDIM]). It allows one to write such code:

from unsafe import gc_alloc, gc_ptr

@struct
class ArrayData[DTYPE]:
    length: i32
    capacity: i32
    items: gc_ptr[DTYPE]

@struct
class Array1d[DTYPE]:

    type ArrData = ArrayData[DTYPE]

    __ll__: gc_ptr[ArrData]

    def __new__(length: i32) -> Self:
        ll = gc_alloc[ArrData](1)
        ll.length = length
        ll.capacity = length
        ll.items = gc_alloc[DTYPE](length)
        return Self.__make__(ll)

def main() -> None:
    a_floats = Array1d[f64](10)
    print(a_floats)

Without the type in type ArrData = ArrayData[DTYPE], it (rightly) fails ("not found in this scope") since this is indeed syntactic sugar for the wrong code

@blue.generic
def Array1d(DTYPE):

    @struct
    class Self:
        # this is a class attribute (not what we want)
        ArrData = ArrayData[DTYPE]

        __ll__: gc_ptr[ArrData]

        def __new__(length: i32) -> Self:
            ll = gc_alloc[ArrData](1)
            ll.length = length
            ll.capacity = length
            ll.items = gc_alloc[DTYPE](length)
            return Self.__make__(ll)

    return Self

With the type, this becomes syntactic sugar for (correct):

@blue.generic
def Array1d(DTYPE):
    # just a local type alias
    ArrData = ArrayData[DTYPE]
    @struct
    class Self:
        __ll__: gc_ptr[ArrData]

        def __new__(length: i32) -> Self:
            ll = gc_alloc[ArrData](1)
            ll.length = length
            ll.capacity = length
            ll.items = gc_alloc[DTYPE](length)
            return Self.__make__(ll)

    return Self

For generic types defined with the syntactic sugar class Array1d[DTYPE]:, we need to be able to define "type aliases" (defined in the outer blue function) and class attributes (defined in the inner struct).

I think the Python ast.TypeAlias syntax (type ArrData = ArrayData[DTYPE]) is adapted for type aliases.

I think this syntax is only useful in generic structs written with the class Array1d[DTYPE]: syntax (maybe also for generic functions written with def func[blue_arg]():, but this is much less useful in this case).

[paugier]

An example of "real" code for which I wish to use that: https://github.com/paugier/spy-demos/blob/loop-fusion/loop-fusion/lib_ndarray.spy#L182. Note

@blue.generic
def LazyAdd(DTYPE, NDIM, LEFT_T, RIGHT_T):
    # a type alias
    A = ndarray[DTYPE, NDIM]

    @struct
    class Self:
        # class attributes
        _DTYPE = DTYPE
        _NDIM = NDIM

        left: LEFT_T
        right: RIGHT_T

and the numbers of A used in this struct. I'd like to write this with class LazyAdd[DTYPE, NDIM, LEFT_T, RIGHT_T]: but then I would need to replace all the A with ndarray[DTYPE, NDIM].

[antocuni]

Without the type in type ArrData = ArrayData[DTYPE], it (rightly) fails ("not found in this scope") since this is indeed syntactic sugar for the wrong code

I think it wrongly fails :).
If you assign a variable inside a class body, you should be able to use that variable inside the class body, as it happens in Python.

There WILL be a difference with python because things like a: int = 4 are not assignments but field declarations, but that's a different story.

I'm also not sure that the type construct is the right choice in SPy. In Python it's present specifically because Python's type system is very ad-hoc and the type checkers might not be able to evaluate type aliases correctly with just =, but in SPy this problem doesn't exist because blue code is evaluated by design.
Also, there is no difference between a type and other compiler time values. Consider this example:

# fictional example
class FixedArray[T, N]:
    bytes_size = T.sizeof() * N
    buf: char_array[bytes_size]

this is exactly the same situation as yours, but bytes_size is not a type and thus type = ... would not work.

Also, don't forget that the class ...[T] form is syntax sugar: I agree that when it works is super nice to use, but at the same time I think it will always be strictly less powerful than the more general @blue.generic form.

Anyway, for the specific case at hand, I think that it should just work out of the box.

[paugier]

I agree that type var = ... is not a very nice syntax since it is not only for types!

If you assign a variable inside a class body, you should be able to use that variable inside the class body, as it happens in Python.

In Python, you cannot use a class variable in the body of a method. Typically, this fails with "NameError: name 'ArrData' is not defined. Did you mean: 'self.ArrData'?":

class ArrayData[DTYPE]:
    length: int
    capacity: int
    items: list[DTYPE]

class Array1d[DTYPE]:
    ArrData = ArrayData[DTYPE]

    def foo(self) -> None:
        print(ArrData)

Do you confirm that you like this to be supported in SPy?

This would be convenient but it is a big change compared to Python (which needs print(self.ArrData)).

It would make sense to only support self.ArrData like in Python but then it does not work for static methods and @blue.metafunc.

[antocuni]

In Python, you cannot use a class variable in the body of a method. Typically, this fails with "NameError: name 'ArrData' is not defined. Did you mean: 'self.ArrData'?":

ah ok, I see what you mean. That's one of the mane python scoping weirdness.
But I think that this problem happens also in CPython/mypy, doesn't it?

class MyList[T]:
    type DATA = list[T]

    def foo(self) -> DATA:
        x: DATA
        return x

This fails with mypy 2.1.0;

❯ uvx mypy a.py 
a.py:2: error: All type parameters should be declared ("T" not declared)  [valid-type]
a.py:5: error: Name "DATA" is not defined  [name-defined]
Found 2 errors in 1 file (checked 1 source file)

I'm undecided between various conflicting solutions here.

(non) solution 1

To declare that we do as in CPython and this pattern just doesn't work.
If you want to you it, you use the @blue.generic unsugared syntax.

solution 2

Abuse the nonlocal statement to mean "declare it in the outer scope". So this would read something like this:

class Array1d[DTYPE]:
    nonlocal ArrData
    ArrData = ArrayData[DTYPE]

    def foo(self) -> None:
        print(ArrData)

Unfortunately nonlocal ArrData = ArrayData[DTYPE] doesn't parse, so we cannot use it.

solution 3

similar to (2), but we modify the parser and we add another modifier similar to what we do already for var/const. Something like this (suggestions for a better name are welcome, maybe even nonlocal):

class Array1d[DTYPE]:
    outer ArrData = ArrayData[DTYPE]

    def foo(self) -> None:
        print(ArrData)

solution 4

declare that class scopes work very differently than CPython:

• bare assignments like x = 1 do NOT becomes class vars but are visible in the inner scopes (what you need)
• if you want a classvar you must explicitly use x: ClassVar[T]
• x: T and x: T = v forms are field declarations, also not visible in the inner scope

This is the bigger departure from CPython semantics, but at the same time it might be the cleanest/closest to existing mypy semantics. But before deciding on this we would need to think deeper because it might interact badly with other cases.

[paugier]

My personal preference goes to abusing nonlocal ArrData = ArrayData[DTYPE] (by modifying the parser like for var/const).

• It's very simple, very consistent with "generic struct with [ARG] are just syntactic sugar for @blue.generic" and thus easy to teach. nonlocal is indeed much better than type. "This variable is nonlocal to the inner red struct and therefore local to the blue.generic function."
• Class variables remain exactly as in Python (I was happy about that when I discovered that they work like that).

Generally, it seems to me that there should be very good reasons to decide to depart from Python semantics. It seems to me that Solution 4 is a big departure compared to how Python and Mypy work, which is not going to be natural for most Python developers. And I don't see the strong advantage of this solution.

Note that currently, blue class variables do not need to be explicitly typed:

@blue.generic
def LazyAdd(DTYPE, NDIM, LEFT_T, RIGHT_T):
    A = ndarray[DTYPE, NDIM]

    @struct
    class Self:
        # not explicitly typed
        _DTYPE = DTYPE
        _NDIM = NDIM