How should pandas-stubs balance typing precision, overload complexity, and runtime behavior for MultiIndex indexing?

[Prem-gosavi45]

Hi pandas maintainers and contributors!

While working on pandas-stubs issue #1018 (MultiIndex column selection with .loc), I encountered an interesting challenge around overload resolution.

A seemingly simple indexing expression can have different return types depending on whether the MultiIndex key is fully specified, partially specified, or contains slices.

For example:

df.loc[:, ("A", "a")]          # Series
df.loc[:, pd.IndexSlice[:, "a"]]  # DataFrame
df.loc[:, pd.IndexSlice["A", :]]  # DataFrame

This raises a broader design question:

How should we model pandas indexing APIs in static typing when the return type depends on the structure and completeness of an index key, especially when Python's type system cannot easily express exclusions between overlapping tuple types?

I'm particularly interested in the following:

1. Should pandas-stubs prioritize precise return-type inference for common indexing patterns, even if that requires more specialized overloads?
2. How should we handle overloads that are semantically distinct at runtime but structurally overlap in Pyright or mypy?
3. Would introducing more specific type aliases for full MultiIndex keys and partial/slice-based keys be a maintainable long-term solution?
4. How do maintainers decide when an overload is worth its complexity versus returning a broader union?
5. Are there plans to establish common typing design principles for complex pandas APIs such as .loc, .iloc, MultiIndex, and GroupBy?

I'd love to understand the maintainers' perspective on the long-term direction of pandas-stubs typing.

Thanks for maintaining this amazing project!

[Gianni300x]

A runtime edge case to include when evaluating an overload for a fully specified key: duplicate column labels can change the result. On pandas 3.0.3 I checked:

import pandas as pd

cols = pd.MultiIndex.from_tuples([("A", "a"), ("A", "a")])
df = pd.DataFrame([[1, 2]], columns=cols)
print(type(df.loc[:, ("A", "a")]).__name__)  # DataFrame

So even a complete tuple alone does not guarantee a Series unless uniqueness is also assumed. The duplicate-label guide discusses how duplicates affect dimensionality. A runtime/type-checker test matrix covering unique and duplicate full keys, partial keys, and slices could make the precision trade-off explicit. I can’t speak for the maintainers’ policy or planned direction.

[Prem-gosavi45]

Thanks, Gianni! That's an important edge case.

You're right that a fully specified MultiIndex key does not necessarily imply a Series, since duplicate column labels can cause .loc to return a DataFrame.

This suggests that key completeness alone is not sufficient to determine the return type. Since column-label uniqueness is generally a runtime property that the type checker cannot guarantee, perhaps a broader return type (Series | DataFrame) is safer for cases where uniqueness cannot be established.

I think a useful next step would be to add a runtime and type-checker test matrix covering:

• Unique full keys
• Duplicate full keys
• Partial keys
• Slice-based keys

I'd be interested to hear whether maintainers prefer this conservative approach or see another maintainable way to model these cases in pandas-stubs.

[niewolski]

I think the best approach would be to prioritize precise type inference for the most common indexing patterns while keeping the overload definitions as simple and maintainable as possible.

For MultiIndex operations, the main challenge is that Python's type system cannot always distinguish between a complete tuple key and a partial tuple key containing slices or other indexing expressions.

A few ideas that could help:

More specific overloads first: Handle common cases, such as fully specified MultiIndex keys returning a Series, before falling back to broader overloads.
Use type aliases: Defining aliases for complete MultiIndex keys and partial indexing keys could improve readability and make the stubs easier to maintain.
Avoid excessive complexity: Trying to cover every possible indexing combination with highly specialized overloads might make the stubs difficult to understand and maintain. A reasonable fallback to a broader return type is sometimes better.
Add regression tests: Testing representative indexing expressions with both mypy and Pyright would help prevent accidental changes in inferred return types.

Personally, I believe the goal should be to provide accurate typing for common real-world usage without trying to model every possible runtime behavior of pandas perfectly.

It would also be useful to establish consistent typing guidelines across APIs such as .loc, .iloc, MultiIndex, and GroupBy, so that future improvements follow the same principles.

I'd be interested to hear how other contributors balance typing precision against long-term maintainability.

[Prem-gosavi45]

Thanks, niewolski! I agree with your approach.

I especially like the idea of prioritizing precise inference for common indexing patterns while keeping overloads simple and maintainable. Type aliases could also help make the distinction between fully specified keys and partial or slice-based keys clearer.

One important consideration is duplicate labels. As Gianni pointed out, even a fully specified MultiIndex key can return a DataFrame when the matching column labels are duplicated. Since uniqueness is generally a runtime property that cannot be reliably established by a static type checker, we should be careful about inferring Series solely from key completeness.

Perhaps a practical balance would be to provide precise overloads where the return type can be determined safely, and use a broader Series | DataFrame return type where runtime conditions such as label uniqueness make the result ambiguous.

I also agree that a representative test matrix covering unique and duplicate full keys, partial keys, and slice-based selections would help keep the behavior consistent across mypy and Pyright.

Thanks for sharing your perspective!

[arpanjot-coding]

One case I'd include in the tests for the overload is duplicate full column labels. On pandas 3.1.1, the exact same key returns a Series with unique labels and a DataFrame when the full label occurs twice.

import pandas as pd

unique = pd.DataFrame(
    [[1, 2]], columns=pd.MultiIndex.from_tuples([("A", "a"), ("A", "b")])
)
duplicate = pd.DataFrame(
    [[1, 2]], columns=pd.MultiIndex.from_tuples([("A", "a"), ("A", "a")])
)

assert isinstance(unique.loc[:, ("A", "a")], pd.Series)
assert isinstance(duplicate.loc[:, ("A", "a")], pd.DataFrame)
assert isinstance(unique.loc[:, [("A", "a")]], pd.DataFrame)

So even a fully specified tuple doesn't tell the checker how many columns match. I'd cover that case alongside the partial and slice keys before making the tuple overload more precise. The list-of-keys form is useful when the caller wants a DataFrame consistently.

[Prem-gosavi45]

Thanks, arpanjot-coding! This is a really useful example.

I agree that duplicate full column labels should be part of the test matrix. It clearly demonstrates that even when a MultiIndex key is fully specified, the return type can depend on runtime label uniqueness.

The distinction between a single tuple key and a list containing a tuple key is also important:

• df.loc[:, ("A", "a")] can return either a Series or a DataFrame, depending on whether the label is unique.
• df.loc[:, [("A", "a")]] consistently returns a DataFrame in your examples.

This reinforces the idea that key completeness alone is not sufficient to determine the return type.

I think a practical approach would be to prioritize precise typing where the result is statically predictable, and use Series | DataFrame where runtime properties such as duplicate labels affect dimensionality.

A test matrix covering unique and duplicate full keys, list-of-tuple selection, partial keys, and slice-based selection would be a good foundation for evaluating the overload design.

Thanks for adding this case!

[arpanjot-coding]

On pandas 3.1.1, the cases with unique and duplicated full column labels behaved as expected. One assertion I'd add is that df.loc[:, [("A", "a")]] stays a DataFrame and retains both matching columns in the duplicate fixture. Both the partial key df.loc[:, "A"] and the slice df.loc[:, pd.IndexSlice[:, "a"]] returned DataFrames in those MultiIndex fixtures.

I'd keep the runtime assertions separate from the type-checker assertions. For a full tuple key, the two fixtures give different runtime types, but I'd expect the same conservative Series | DataFrame type for both. The partial string key has a similar distinction. df.loc[:, "A"] gives a DataFrame with these MultiIndex columns and a Series on a flat frame with a unique column named "A". A test for each would guard against making the general string-key overload too narrow.