Add mypy annotations for _api_calls.py (#1257)

* Add mypy annotations for _api_calls.py

This commit adds mypy annotations for _api_calls.py to make
sure that we have no untyped classes and functions.

* Include Pieters suggestions

Author: mfeurer

.pre-commit-config.yaml

@@ -19,6 +19,14 @@ repos:
         additional_dependencies:
           - types-requests
           - types-python-dateutil
+      - id: mypy
+        name: mypy top-level-functions
+        files: openml/_api_calls.py
+        additional_dependencies:
+          - types-requests
+          - types-python-dateutil
+        args: [ --disallow-untyped-defs, --disallow-any-generics,
+                --disallow-any-explicit, --implicit-optional ]
   - repo: https://github.com/pycqa/flake8
     rev: 6.0.0
     hooks:

openml/_api_calls.py

@@ -11,7 +11,7 @@
 import xml
 import xmltodict
 from urllib3 import ProxyManager
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Tuple, Union
 import zipfile
 
 import minio
@@ -24,6 +24,9 @@
     OpenMLHashException,
 )
 
+DATA_TYPE = Dict[str, Union[str, int]]
+FILE_ELEMENTS_TYPE = Dict[str, Union[str, Tuple[str, str]]]
+
 
 def resolve_env_proxies(url: str) -> Optional[str]:
     """Attempt to find a suitable proxy for this url.
@@ -54,7 +57,12 @@ def _create_url_from_endpoint(endpoint: str) -> str:
     return url.replace("=", "%3d")
 
 
-def _perform_api_call(call, request_method, data=None, file_elements=None):
+def _perform_api_call(
+    call: str,
+    request_method: str,
+    data: Optional[DATA_TYPE] = None,
+    file_elements: Optional[FILE_ELEMENTS_TYPE] = None,
+) -> str:
     """
     Perform an API call at the OpenML server.
 
@@ -76,8 +84,6 @@ def _perform_api_call(call, request_method, data=None, file_elements=None):
 
     Returns
     -------
-    return_code : int
-        HTTP return code
     return_value : str
         Return value of the OpenML server
     """
@@ -257,7 +263,7 @@ def _download_text_file(
         return None
 
 
-def _file_id_to_url(file_id, filename=None):
+def _file_id_to_url(file_id: str, filename: Optional[str] = None) -> str:
     """
     Presents the URL how to download a given file id
     filename is optional
@@ -269,7 +275,9 @@ def _file_id_to_url(file_id, filename=None):
     return url
 
 
-def _read_url_files(url, data=None, file_elements=None):
+def _read_url_files(
+    url: str, data: Optional[DATA_TYPE] = None, file_elements: Optional[FILE_ELEMENTS_TYPE] = None
+) -> requests.Response:
     """do a post request to url with data
     and sending file_elements as files"""
 
@@ -288,7 +296,12 @@ def _read_url_files(url, data=None, file_elements=None):
     return response
 
 
-def __read_url(url, request_method, data=None, md5_checksum=None):
+def __read_url(
+    url: str,
+    request_method: str,
+    data: Optional[DATA_TYPE] = None,
+    md5_checksum: Optional[str] = None,
+) -> requests.Response:
     data = {} if data is None else data
     if config.apikey:
         data["api_key"] = config.apikey
@@ -306,10 +319,16 @@ def __is_checksum_equal(downloaded_file_binary: bytes, md5_checksum: Optional[st
     return md5_checksum == md5_checksum_download
 
 
-def _send_request(request_method, url, data, files=None, md5_checksum=None):
+def _send_request(
+    request_method: str,
+    url: str,
+    data: DATA_TYPE,
+    files: Optional[FILE_ELEMENTS_TYPE] = None,
+    md5_checksum: Optional[str] = None,
+) -> requests.Response:
     n_retries = max(1, config.connection_n_retries)
 
-    response = None
+    response: requests.Response
     with requests.Session() as session:
         # Start at one to have a non-zero multiplier for the sleep
         for retry_counter in range(1, n_retries + 1):
@@ -380,12 +399,12 @@ def human(n: int) -> float:
 
                     delay = {"human": human, "robot": robot}[config.retry_policy](retry_counter)
                     time.sleep(delay)
-    if response is None:
-        raise ValueError("This should never happen!")
     return response
 
 
-def __check_response(response, url, file_elements):
+def __check_response(
+    response: requests.Response, url: str, file_elements: Optional[FILE_ELEMENTS_TYPE]
+) -> None:
     if response.status_code != 200:
         raise __parse_server_exception(response, url, file_elements=file_elements)
     elif (
@@ -397,7 +416,7 @@ def __check_response(response, url, file_elements):
 def __parse_server_exception(
     response: requests.Response,
     url: str,
-    file_elements: Dict,
+    file_elements: Optional[FILE_ELEMENTS_TYPE],
 ) -> OpenMLServerError:
     if response.status_code == 414:
         raise OpenMLServerError("URI too long! ({})".format(url))

openml/datasets/functions.py

@@ -4,7 +4,7 @@
 import logging
 import os
 from pyexpat import ExpatError
-from typing import List, Dict, Union, Optional, cast
+from typing import List, Dict, Optional, Union, cast
 import warnings
 
 import numpy as np
@@ -867,7 +867,7 @@ def edit_dataset(
         raise TypeError("`data_id` must be of type `int`, not {}.".format(type(data_id)))
 
     # compose data edit parameters as xml
-    form_data = {"data_id": data_id}
+    form_data = {"data_id": data_id}  # type: openml._api_calls.DATA_TYPE
     xml = OrderedDict()  # type: 'OrderedDict[str, OrderedDict]'
     xml["oml:data_edit_parameters"] = OrderedDict()
     xml["oml:data_edit_parameters"]["@xmlns:oml"] = "http://openml.org/openml"
@@ -888,7 +888,9 @@ def edit_dataset(
         if not xml["oml:data_edit_parameters"][k]:
             del xml["oml:data_edit_parameters"][k]
 
-    file_elements = {"edit_parameters": ("description.xml", xmltodict.unparse(xml))}
+    file_elements = {
+        "edit_parameters": ("description.xml", xmltodict.unparse(xml))
+    }  # type: openml._api_calls.FILE_ELEMENTS_TYPE
     result_xml = openml._api_calls._perform_api_call(
         "data/edit", "post", data=form_data, file_elements=file_elements
     )
@@ -929,7 +931,7 @@ def fork_dataset(data_id: int) -> int:
     if not isinstance(data_id, int):
         raise TypeError("`data_id` must be of type `int`, not {}.".format(type(data_id)))
     # compose data fork parameters
-    form_data = {"data_id": data_id}
+    form_data = {"data_id": data_id}  # type: openml._api_calls.DATA_TYPE
     result_xml = openml._api_calls._perform_api_call("data/fork", "post", data=form_data)
     result = xmltodict.parse(result_xml)
     data_id = result["oml:data_fork"]["oml:id"]
@@ -949,7 +951,7 @@ def _topic_add_dataset(data_id: int, topic: str):
     """
     if not isinstance(data_id, int):
         raise TypeError("`data_id` must be of type `int`, not {}.".format(type(data_id)))
-    form_data = {"data_id": data_id, "topic": topic}
+    form_data = {"data_id": data_id, "topic": topic}  # type: openml._api_calls.DATA_TYPE
     result_xml = openml._api_calls._perform_api_call("data/topicadd", "post", data=form_data)
     result = xmltodict.parse(result_xml)
     data_id = result["oml:data_topic"]["oml:id"]
@@ -970,7 +972,7 @@ def _topic_delete_dataset(data_id: int, topic: str):
     """
     if not isinstance(data_id, int):
         raise TypeError("`data_id` must be of type `int`, not {}.".format(type(data_id)))
-    form_data = {"data_id": data_id, "topic": topic}
+    form_data = {"data_id": data_id, "topic": topic}  # type: openml._api_calls.DATA_TYPE
     result_xml = openml._api_calls._perform_api_call("data/topicdelete", "post", data=form_data)
     result = xmltodict.parse(result_xml)
     data_id = result["oml:data_topic"]["oml:id"]

openml/setups/functions.py

@@ -49,7 +49,9 @@ def setup_exists(flow) -> int:
 
     openml_param_settings = flow.extension.obtain_parameter_values(flow)
     description = xmltodict.unparse(_to_dict(flow.flow_id, openml_param_settings), pretty=True)
-    file_elements = {"description": ("description.arff", description)}
+    file_elements = {
+        "description": ("description.arff", description)
+    }  # type: openml._api_calls.FILE_ELEMENTS_TYPE
     result = openml._api_calls._perform_api_call(
         "/setup/exists/", "post", file_elements=file_elements
     )

openml/study/functions.py

@@ -277,7 +277,7 @@ def update_study_status(study_id: int, status: str) -> None:
     legal_status = {"active", "deactivated"}
     if status not in legal_status:
         raise ValueError("Illegal status value. " "Legal values: %s" % legal_status)
-    data = {"study_id": study_id, "status": status}
+    data = {"study_id": study_id, "status": status}  # type: openml._api_calls.DATA_TYPE
     result_xml = openml._api_calls._perform_api_call("study/status/update", "post", data=data)
     result = xmltodict.parse(result_xml)
     server_study_id = result["oml:study_status_update"]["oml:id"]
@@ -357,8 +357,10 @@ def attach_to_study(study_id: int, run_ids: List[int]) -> int:
 
     # Interestingly, there's no need to tell the server about the entity type, it knows by itself
     uri = "study/%d/attach" % study_id
-    post_variables = {"ids": ",".join(str(x) for x in run_ids)}
-    result_xml = openml._api_calls._perform_api_call(uri, "post", post_variables)
+    post_variables = {"ids": ",".join(str(x) for x in run_ids)}  # type: openml._api_calls.DATA_TYPE
+    result_xml = openml._api_calls._perform_api_call(
+        call=uri, request_method="post", data=post_variables
+    )
     result = xmltodict.parse(result_xml)["oml:study_attach"]
     return int(result["oml:linked_entities"])
 
@@ -400,8 +402,10 @@ def detach_from_study(study_id: int, run_ids: List[int]) -> int:
 
     # Interestingly, there's no need to tell the server about the entity type, it knows by itself
     uri = "study/%d/detach" % study_id
-    post_variables = {"ids": ",".join(str(x) for x in run_ids)}
-    result_xml = openml._api_calls._perform_api_call(uri, "post", post_variables)
+    post_variables = {"ids": ",".join(str(x) for x in run_ids)}  # type: openml._api_calls.DATA_TYPE
+    result_xml = openml._api_calls._perform_api_call(
+        call=uri, request_method="post", data=post_variables
+    )
     result = xmltodict.parse(result_xml)["oml:study_detach"]
     return int(result["oml:linked_entities"])