Add type hints to public API functions

[goransh-buh]

Feature Request

Summary

Adding type hints to public API functions would significantly improve developer experience and enable better tooling support.

Motivation

• IDEs can provide better autocomplete and inline documentation
• Static type checkers like mypy can catch bugs before runtime
• New contributors can understand the expected inputs/outputs faster
• Type hints serve as machine-checked documentation

Proposed Change

Add type annotations to all public functions:

# Before:
def process(data, config=None):
    ...

# After:
from typing import Optional, Dict, Any

def process(data: Dict[str, Any], config: Optional[Dict] = None) -> bool:
    ...

Additional Context

This is a non-breaking change and can be done incrementally, starting with the most-used public functions.

[smqd19]

Adding type hints to public API functions is a great idea for improving developer experience and tooling support. In our experience at Reallytics.ai, we've found that type hints significantly enhance the productivity of our development team and reduce the number of runtime errors. We've used mypy for static type checking and have seen a notable decrease in bugs that could have been caught during development. Additionally, type hints serve as a form of machine-checked documentation, making it easier for new contributors to understand the expected inputs and outputs.

Here's a concrete suggestion: Start with the most frequently used public functions, as suggested. This incremental approach will allow you to gauge the impact and make any necessary adjustments before rolling out the change across the entire codebase. It's also a good idea to document the changes in the README or inline comments to provide context for new contributors.

In terms of tools, we've found that mypy is a powerful static type checker that integrates well with popular IDEs like PyCharm and VSCode. It can be configured to check the entire codebase or specific files, making it easy to maintain and update as the project evolves.

[AyushKashyapII]

hey @goransh-buh I took a stab at this. I've made a focused start on the aggregation API in Ray Data ,specifically aggregate.py and grouped_data.py.

The fixes include:

AggregateFn.init: accumulate_row and accumulate_block were typed as Callable[...] = None but None was not in the union — wrapped both with Optional[...]
GroupedData.sum, min, max, mean, std: on: Union[str, List[str]] = None → on: Optional[Union[str, List[str]]] = None (same issue)
GroupedData.map_groups: same fix for the compute parameter
Added -> None to all public init methods across both files (18 in aggregate.py, 1 in grouped_data.py)
Added missing -> Dataset return type to _aggregate_on
These are mypy-catchable errors, not just missing annotations, so they have real correctness value. PR is up — happy to expand to other files in follow-up PRs if the maintainers find this approach useful!

Thanks