Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/array_class.py (Python)

Code before:
import math
import dataclasses
from functools import partial

import gstaichi as ti
import numpy as np

import genesis as gs


if not gs._initialized:
    gs.raise_exception("Genesis hasn't been initialized. Did you call `gs.init()`?")


V_ANNOTATION = ti.types.ndarray() if gs.use_ndarray else ti.template
V = ti.ndarray if gs.use_ndarray else ti.field
V_VEC = ti.Vector.ndarray if gs.use_ndarray else ti.Vector.field
V_MAT = ti.Matrix.ndarray if gs.use_ndarray else ti.Matrix.field

DATA_ORIENTED = partial(dataclasses.dataclass, frozen=True) if gs.use_ndarray else ti.data_oriented


def _shape(shape, is_on):
    return shape if is_on else ()


class AutoInitMeta(type):
    def __new__(cls, name, bases, namespace):
        field_names = namespace["__annotations__"].keys()

        def __init__(self, *args, **kwargs):
            assigned = {}

            # Assign positional arguments
            if len(args) > len(field_names):
                raise TypeError(f"{name}() takes {len(field_names)} positional arguments but {len(args)} were given")
            for field, value in zip(field_names, args):
                assigned[field] = value

            # Assign keyword arguments
            for key, value in kwargs.items():
                if key not in field_names:
                    raise TypeError(f"{name}() got unexpected keyword argument '{key}'")
                if key in assigned:
                    raise TypeError(f"{name}() got multiple values for argument '{key}'")
                assigned[key] = value

            # Check for missing arguments

Review comment (suggestion):
I would be better to have a more explicit name if you want to keep this pattern. Like `optional_shape` or `cond_shape` or `maybe_shape`.

Code after:
import math
import dataclasses
from functools import partial

import gstaichi as ti
import numpy as np

import genesis as gs


if not gs._initialized:
    gs.raise_exception("Genesis hasn't been initialized. Did you call `gs.init()`?")


V_ANNOTATION = ti.types.ndarray() if gs.use_ndarray else ti.template
V = ti.ndarray if gs.use_ndarray else ti.field
V_VEC = ti.Vector.ndarray if gs.use_ndarray else ti.Vector.field
V_MAT = ti.Matrix.ndarray if gs.use_ndarray else ti.Matrix.field

DATA_ORIENTED = partial(dataclasses.dataclass, frozen=True) if gs.use_ndarray else ti.data_oriented


def maybe_shape(shape, is_on):
    return shape if is_on else ()


class AutoInitMeta(type):
    def __new__(cls, name, bases, namespace):
        field_names = namespace["__annotations__"].keys()

        def __init__(self, *args, **kwargs):
            assigned = {}

            # Assign positional arguments
            if len(args) > len(field_names):
                raise TypeError(f"{name}() takes {len(field_names)} positional arguments but {len(args)} were given")
            for field, value in zip(field_names, args):
                assigned[field] = value

            # Assign keyword arguments
            for key, value in kwargs.items():
                if key not in field_names:
                    raise TypeError(f"{name}() got unexpected keyword argument '{key}'")
                if key in assigned:
                    raise TypeError(f"{name}() got multiple values for argument '{key}'")
                assigned[key] = value

            # Check for missing arguments