[Feature Request] Consider `PyStructSequence` types as internal node types

Required prerequisites

[X] I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Motivation

Currently, only tuple and subclasses of namedtuple are considered node types. Other subclasses of tuple should be manually registered as node type. Otherwise, they are opaque leaf type.

In Python documentation Struct Sequence Objects:

Struct sequence objects are the C equivalent of namedtuple() objects, i.e. a sequence whose items can also be accessed through attributes. To create a struct sequence, you first have to create a specific struct sequence type.

PyStructSequence types are also subclasses of tuple. They are widely used in Python standard libraries (e.g, time.struct_time) and other packages (e.g., torch.return_types.*).

It would be nice if we consider PyStructSequence types as internal node types by default.

Ref:

pytorch/pytorch#75982
What’s the best way to check if a type is PyStructSequence

Solution

Add special handling of PyStructSequence subclasses in C extensions as we did for namedtuple.

A snippet I posted in What’s the best way to check if a type is PyStructSequence:

import inspect

def is_namedtuple(obj: object | type) -> bool:
    cls = obj if inspect.isclass(cls) else type(obj)
    return (
        issubclass(cls, tuple)
        and isinstance(getattr(cls, '_fields', None), tuple)
        and all(isinstance(field, str) for field in cls._fields)
    )

def is_structseq(obj: object | type) -> bool:
    cls = obj if inspect.isclass(cls) else type(obj)
    if (
        cls.__base__ is tuple
        and isinstance(getattr(cls, 'n_sequence_fields', None), int)
        and isinstance(getattr(cls, 'n_fields', None), int)
        and isinstance(getattr(cls, 'n_unnamed_fields', None), int)
    ):
        try:

            class subcls(cls):
                pass

        except (
            TypeError,       # CPython
            AssertionError,  # PyPy
        ):
            return True

    return False

>>> is_structseq(torch.return_types.max)
True
>>> is_structseq(time.struct_time)
True

Alternatives

Keep current API unchanged. Users need to register the PyStructSequence types manually.

Additional context

No response

metaopt / optree