diff --git a/packages/typemap/src/typemap/type_eval/_eval_operators.py b/packages/typemap/src/typemap/type_eval/_eval_operators.py index 73db59d..9009b69 100644 --- a/packages/typemap/src/typemap/type_eval/_eval_operators.py +++ b/packages/typemap/src/typemap/type_eval/_eval_operators.py @@ -36,6 +36,7 @@ IsAssignable, IsEquivalent, Iter, + KeyOf, Length, Lowercase, Member, @@ -1138,6 +1139,27 @@ def _eval_Length(tp, *, ctx) -> typing.Any: raise TypeError(f"Invalid type argument to Length: {tp} is not a tuple") +@type_eval.register_evaluator(KeyOf) +@_lift_over_unions +def _eval_KeyOf(tp, *, ctx): + """Evaluate KeyOf[T] to get all member names as a tuple of Literals.""" + tp = _eval_types(tp, ctx) + hints = get_annotated_type_hints( + tp, include_extras=True, attrs_only=True, ctx=ctx + ) + + if not hints: + return typing.Literal[()] + + # Extract member names and create tuple of Literal types + names = [] + for name in hints: + names.append(typing.Literal[name]) + + # Return as tuple of Literal types (use unpacking to make it hashable) + return tuple[*names] # type: ignore[return-value] + + @type_eval.register_evaluator(Slice) @_lift_over_unions def _eval_Slice(tp, start, end, *, ctx): diff --git a/packages/typemap/src/typemap/typing.py b/packages/typemap/src/typemap/typing.py index df3db5a..880835c 100644 --- a/packages/typemap/src/typemap/typing.py +++ b/packages/typemap/src/typemap/typing.py @@ -262,6 +262,19 @@ class Length[S: tuple]: pass +class KeyOf[T]: + """Return all member names as a tuple. + + Usage: + type Keys = KeyOf[User] + # Returns tuple[Literal['name'], Literal['email'], Literal['age']] + + This is similar to TypeScript's keyof operator. + """ + + pass + + class Slice[S: str | tuple, Start: int | None, End: int | None]( _TupleLikeOperator ): diff --git a/packages/typemap/tests/test_type_eval.py b/packages/typemap/tests/test_type_eval.py index f557e4d..89439b4 100644 --- a/packages/typemap/tests/test_type_eval.py +++ b/packages/typemap/tests/test_type_eval.py @@ -37,6 +37,7 @@ GetAnnotations, IsAssignable, Iter, + KeyOf, Length, IsEquivalent, Member, @@ -2648,3 +2649,77 @@ def test_raise_error_with_literal_types(): eval_typing( RaiseError[Literal["Shape mismatch"], Literal[4], Literal[3]] ) + + +############## +# KeyOf tests + + +def test_keyof_basic(): + """Test KeyOf returns tuple of Literal member names.""" + + class User: + name: str + age: int + email: str + + result = eval_typing(KeyOf[User]) + assert ( + result + == tuple[ + Literal["name"], + Literal["age"], + Literal["email"], + ] + ) + + +def test_keyof_single_field(): + """Test KeyOf with a single field class.""" + + class Single: + value: int + + result = eval_typing(KeyOf[Single]) + assert result == tuple[Literal["value"]] + + +def test_keyof_empty_class(): + """Test KeyOf with a class with no fields.""" + + class Empty: + pass + + result = eval_typing(KeyOf[Empty]) + assert result == Literal[()] + + +def test_keyof_with_methods(): + """Test KeyOf ignores methods and only returns field names.""" + + class WithMethods: + name: str + + def method(self) -> None: ... + + result = eval_typing(KeyOf[WithMethods]) + assert result == tuple[Literal["name"]] + + +def test_keyof_with_inheritance(): + """Test KeyOf includes inherited fields.""" + + class Base: + id: int + + class Child(Base): + name: str + + result = eval_typing(KeyOf[Child]) + assert ( + result + == tuple[ + Literal["id"], + Literal["name"], + ] + )