# -------------------------------------- ARRANGE ------------------------------------- #

import polars as pl
import dataframely as dy

class Schema(dy.Schema):
    a = dy.Int64(primary_key=True)

class ParentCollection(dy.Collection):
    s1: dy.LazyFrame[Schema]

class ChildCollection(ParentCollection):
    s2: dy.LazyFrame[Schema]

    @dy.filter()
    def rule(self) -> dy.LazyFrame[Schema]:
        return self.s1.join(self.s2, on="a", how="inner")

valid, failures = ChildCollection.filter(
    {"s1": pl.DataFrame({"a": [1, 2, 3]}), "s2": pl.DataFrame({"a": [2, 3, 4]})},
    cast=True,
)

# --------------------------------------- ISSUE -------------------------------------- #

a = failures[
    "s1"
]  # Mypy: "s1" is not a valid TypedDict key; expected one of ("s2") Mypymisc
b = failures["s2"]