Update documentation about data source signatures.

PiperOrigin-RevId: 628361784

Author: marcenacp

docs/data_source.ipynb

@@ -172,21 +172,18 @@
         "following protocol:\n",
         "\n",
         "```python\n",
+        "from typing import SupportsIndex\n",
+        "\n",
         "class RandomAccessDataSource(Protocol):\n",
         "  \"\"\"Interface for datasources where storage supports efficient random access.\"\"\"\n",
         "\n",
         "  def __len__(self) -> int:\n",
         "    \"\"\"Number of records in the dataset.\"\"\"\n",
         "\n",
-        "  def __getitem__(self, record_key: int) -> Sequence[Any]:\n",
-        "    \"\"\"Retrieves records for the given record_keys.\"\"\"\n",
+        "  def __getitem__(self, key: SupportsIndex) -> Any:\n",
+        "    \"\"\"Retrieves the record for the given key.\"\"\"\n",
         "```\n",
         "\n",
-        "**Warning**: the API is still under active development. Notably, at this point,\n",
-        "`__getitem__` must support both `int` and `list[int]` in input. In the future,\n",
-        "it will probably only support `int` as per\n",
-        "[the standard](https://docs.python.org/3/reference/datamodel.html#object.__getitem__).\n",
-        "\n",
         "The underlying file format needs to support efficient random access. At the\n",
         "moment, TFDS relies on [`array_record`](https://github.com/google/array_record).\n",
         "\n",