feat(checker): handle assignment to multiple columns

This commit is contained in:
2026-07-02 23:29:10 +02:00
parent b14f46d405
commit 0288a05901

View File

@@ -47,12 +47,41 @@ class FrameManager:
return self.assign_column(reporter, location, frame, name, value_type)
case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
isinstance(idx, str) for idx in indices
isinstance(index.value, str) for index in indices
):
raise NotImplementedError
names: list[str] = [cast(str, index.value) for index in indices]
if not isinstance(value_type, TupleType):
reporter.error(
location,
f"Cannot assign {type} to dataframe columns. Must be a tuple of columns",
)
return UnknownType()
if len(names) != len(value_type.items):
reporter.error(
location,
f"Wrong number of columns. Cannot assign {len(value_type.items)} to {len(names)} targets",
)
return UnknownType()
new_frame: Type = frame
for name, value in zip(names, value_type.items):
new_frame = self.assign_column(
reporter,
location,
new_frame,
name,
value,
)
if not isinstance(new_frame, DataFrameType):
return new_frame
return new_frame
case _:
reporter.error(location, f"Invalid index type {index} on {frame}")
reporter.error(
location, f"Invalid index type {index} on {frame} (assignment)"
)
return UnknownType()
def assign_column(
@@ -100,7 +129,9 @@ class FrameManager:
return TupleType(items=tuple(columns))
case _:
reporter.error(location, f"Invalid index type {index} on {frame}")
reporter.error(
location, f"Invalid index type {index} on {frame} (access)"
)
return UnknownType()
def groupby_get(