Typing of Dataset.__getitem__ could be improved.
MEMBER
Hi ! Thanks for the suggestion, I didn't know about this decorator.

If you are interesting in contributing, feel free to open a pull request to add the overload methods for each typing combination :) To assign you to this issue, you can comment `#self-assign` in this thread.

`Dataset.__getitem__` is defined right here: https://github.com/huggingface/datasets/blob/e6f1352fe19679de897f3d962e616936a17094f5/src/datasets/arrow_dataset.py#L1840
2021-11-29T17:23:32Z
https://github.com/huggingface/datasets/issues/3337#issuecomment-981849174
https://api.github.com/repos/huggingface/datasets/issues/3337
IC_kwDODunzps46hdRW
https://api.github.com/repos/huggingface/datasets/issues/comments/981849174/reactions
2021-11-29T17:23:32Z
https://api.github.com/repos/huggingface/datasets/issues/comments/981849174
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
#self-assign
2021-11-29T17:48:17Z
https://github.com/huggingface/datasets/issues/3337#issuecomment-981868857
https://api.github.com/repos/huggingface/datasets/issues/3337
IC_kwDODunzps46hiE5
https://api.github.com/repos/huggingface/datasets/issues/comments/981868857/reactions
2021-11-29T17:48:17Z
https://api.github.com/repos/huggingface/datasets/issues/comments/981868857
https://avatars.githubusercontent.com/u/8976546?v=4
https://api.github.com/users/Dref360/events{/privacy}
https://api.github.com/users/Dref360/followers
https://api.github.com/users/Dref360/following{/other_user}
https://api.github.com/users/Dref360/gists{/gist_id}
https://github.com/Dref360
Dref360
MDQ6VXNlcjg5NzY1NDY=
https://api.github.com/users/Dref360/orgs
https://api.github.com/users/Dref360/received_events
https://api.github.com/users/Dref360/repos
https://api.github.com/users/Dref360/starred{/owner}{/repo}
https://api.github.com/users/Dref360/subscriptions
User
https://api.github.com/users/Dref360
public
## Describe the bug

The newly added typing for Dataset.__getitem__ is Union[Dict, List]. This makes tools like mypy a bit awkward to use as we need to check the type manually. We could use type overloading to make this easier. [Documentation](https://docs.python.org/3/library/typing.html#typing.overload)

## Steps to reproduce the bug
Let's have a file `test.py`

```python
from typing import List, Dict, Any

from datasets import Dataset

ds = Dataset.from_dict({
    'a': [1,2,3],
    'b': ["1", "2", "3"]
})

one_colum: List[str] = ds['a']
some_index: Dict[Any, Any] = ds[1]
```

## Expected results

Running `mypy test.py` should not give any error.


## Actual results

```
test.py:10: error: Incompatible types in assignment (expression has type "Union[Dict[Any, Any], List[Any]]", variable has type "List[str]")
test.py:11: error: Incompatible types in assignment (expression has type "Union[Dict[Any, Any], List[Any]]", variable has type "Dict[Any, Any]")
Found 2 errors in 1 file (checked 1 source file)
```


## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.13.3
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.8.8
- PyArrow version: 6.0.1
