Final run, various fn fixes. DocFmt now handles decorated fns

cognite/client/_api/data_modeling/_data_modeling_executor.py

@@ -15,7 +15,7 @@ def get_data_modeling_executor() -> TaskExecutor:
     Thus, we use a dedicated executor for these endpoints to match the backend.
 
     Returns:
-        The data modeling executor.
+        TaskExecutor: The data modeling executor.
     """
     global _THREAD_POOL_EXECUTOR_SINGLETON
 

cognite/client/_api/functions.py

@@ -572,7 +572,7 @@ def _extract_requirements_from_file(file_name: str) -> List[str]:
         file_name (str): name of the file to parse
 
     Returns:
-        (list[str]): returns a list of library records
+        List[str]: returns a list of library records
     """
     requirements: List[str] = []
     with open(file_name, "r+") as f:
@@ -590,7 +590,7 @@ def _extract_requirements_from_doc_string(docstr: str) -> Optional[List[str]]:
         docstr (str): the docstring to extract requirements from
 
     Returns:
-        (list[str] | None): returns a list of library records if requirements are defined in the docstring, else None
+        Optional[List[str]]: returns a list of library records if requirements are defined in the docstring, else None
     """
     substr_start, substr_end = None, None
 
@@ -612,7 +612,7 @@ def _validate_and_parse_requirements(requirements: List[str]) -> List[str]:
     """Validates the requirement specifications
 
     Args:
-        requirements (list[str]): list of requirement specifications
+        requirements (List[str]): list of requirement specifications
     Raises:
         ValueError: if validation of requirements fails
     Returns:
@@ -636,7 +636,6 @@ def _get_fn_docstring_requirements(fn: Callable) -> List[str]:
 
     Args:
         fn (Callable): the function to read requirements from
-        file_path (str): Path of file to write requirements to
 
     Returns:
         List[str]: A (possibly empty) list of requirements.

cognite/client/testing.py

@@ -159,7 +159,7 @@ def monkeypatch_cognite_client() -> Iterator[CogniteClientMock]:
     Will patch all clients and replace them with specced MagicMock objects.
 
     Yields:
-        CogniteClientMock: The mock with which the CogniteClient has been replaced
+        Iterator[CogniteClientMock]: The mock with which the CogniteClient has been replaced
 
     Examples:
 

cognite/client/utils/_concurrency.py

@@ -251,7 +251,14 @@ def execute_tasks(
     """
     Will use a default executor if one is not passed explicitly. The default executor type uses a thread pool but can
     be changed using ExecutorSettings.executor_type.
-    """
+
+    Args:
+        func (Callable[..., T_Result]): No description.
+        tasks (Union[Sequence[Tuple], List[Dict]]): No description.
+        max_workers (int): No description.
+        executor (Optional[TaskExecutor]): No description.
+    Returns:
+        TasksSummary: No description."""
     executor = executor or get_executor(max_workers)
     futures = []
     for task in tasks:

cognite/client/utils/_graph.py

@@ -8,7 +8,12 @@ def find_all_cycles_with_elements(has_cycles: Set[str], edges: Dict) -> List[Lis
 
     Raises:
         KeyError: No loop found or edge does not exist.
-    """
+
+    Args:
+        has_cycles (Set[str]): No description.
+        edges (Dict): No description.
+    Returns:
+        List[List[str]]: No description."""
     cycles = []
     skip: Set[str] = set()
     while has_cycles:
@@ -26,7 +31,13 @@ def find_extended_cycle(slow: str, edges: Dict, skip: Set[str]) -> Tuple[Set[str
 
     Raises:
         KeyError: No loop found or edge does not exist.
-    """
+
+    Args:
+        slow (str): No description.
+        edges (Dict): No description.
+        skip (Set[str]): No description.
+    Returns:
+        Tuple[Set[str], List[str]]: No description."""
     all_elements = {slow}
     fast = edges[slow]
     while slow != fast:

cognite/client/utils/_text.py

@@ -45,7 +45,11 @@ def convert_all_keys_to_camel_case_recursive(dct: dict[str, Any]) -> dict[str, A
     """Converts all the dictionary keys from snake to camel cases included nested objects.
     >>> convert_all_keys_to_camel_case_recursive({"my_key": {"my_key": 1}})
     {'myKey': {'myKey': 1}}
-    """
+
+    Args:
+        dct (dict[str, Any]): No description.
+    Returns:
+        dict[str, Any]: No description."""
     return {
         to_camel_case(k): (convert_all_keys_to_camel_case_recursive(v) if isinstance(v, dict) else v)
         for k, v in dct.items()
@@ -56,7 +60,12 @@ def convert_all_keys_recursive(dct: dict[str, Any], camel_case: bool = False) ->
     """Converts all the dictionary keys from snake to camel cases included nested objects.
     >>> convert_all_keys_recursive({"my_key": {"my_key": 1}}, camel_case=True)
     {'myKey': {'myKey': 1}}
-    """
+
+    Args:
+        dct (dict[str, Any]): No description.
+        camel_case (bool): No description.
+    Returns:
+        dict[str, Any]: No description."""
     return {
         (to_camel_case(k) if camel_case else k): (
             convert_all_keys_recursive(v, camel_case) if isinstance(v, dict) else v

cognite/client/utils/_time.py

@@ -68,7 +68,7 @@ def datetime_to_ms(dt: datetime) -> int:
         dt (datetime): Naive or aware datetime object. Naive datetimes are interpreted as local time.
 
     Returns:
-        ms: Milliseconds since epoch (negative for time prior to 1970-01-01)
+        int: Milliseconds since epoch (negative for time prior to 1970-01-01)
     """
     try:
         return int(1000 * dt.timestamp())
@@ -398,12 +398,12 @@ def align_large_granularity(start: datetime, end: datetime, granularity: str) ->
     This is done to get consistent behavior with the Cognite Datapoints API.
 
     Args:
-        start: Start time
-        end: End time
-        granularity: The large granularity, day|week|month|quarter|year.
+        start (datetime): Start time
+        end (datetime): End time
+        granularity (str): The large granularity, day|week|month|quarter|year.
 
     Returns:
-        start and end aligned with granularity
+        tuple[datetime, datetime]: start and end aligned with granularity
     """
     multiplier, unit = get_granularity_multiplier_and_unit(granularity)
     # Can be replaced by a single dispatch pattern, but kept more explicit for readability.
@@ -539,7 +539,14 @@ def pandas_date_range_tz(start: datetime, end: datetime, freq: str, inclusive: s
     This function overcomes that limitation.
 
     Assumes that start and end have the same timezone.
-    """
+
+    Args:
+        start (datetime): No description.
+        end (datetime): No description.
+        freq (str): No description.
+        inclusive (str): No description.
+    Returns:
+        pandas.DatetimeIndex: No description."""
     pd = cast(Any, local_import("pandas"))
     # There is a bug in date_range which makes it fail to handle ambiguous timestamps when you use time zone aware
     # datetimes. This is a workaround by passing the time zone as an argument to the function.
@@ -573,7 +580,12 @@ def _timezones_are_equal(start_tz: tzinfo, end_tz: tzinfo) -> bool:
     Note:
         We do not consider timezones with different keys, but equal fixed offsets from UTC to be equal. An example
         would be Zulu Time (which is +00:00 ahead of UTC) and UTC.
-    """
+
+    Args:
+        start_tz (tzinfo): No description.
+        end_tz (tzinfo): No description.
+    Returns:
+        bool: No description."""
     if start_tz is end_tz:
         return True
     ZoneInfo, ZoneInfoNotFoundError = import_zoneinfo(), _import_zoneinfo_not_found_error()
@@ -630,7 +642,12 @@ def _unit_in_days(unit: str, ceil: bool = True) -> float:
     **Caveat** Should not be used for precise calculations, as month, quarter, and year
     do not have a precise timespan in days. Instead, the ceil argument is used to select between
     the maximum and minimum length of a year, quarter, and month.
-    """
+
+    Args:
+        unit (str): No description.
+        ceil (bool): No description.
+    Returns:
+        float: No description."""
     if unit in {"w", "d", "h", "m", "s"}:
         unit = GRANULARITY_IN_TIMEDELTA_UNIT[unit]
         arg = {unit: 1}
@@ -650,12 +667,11 @@ def in_timedelta(granularity: str, ceil: bool = True) -> timedelta:
     Converts the granularity to a timedelta.
 
     Args:
-        granularity: The granularity.
-        ceil: In the case the unit is month, quarter or year. Ceil = True will use 31, 92, 366 days for these
-              timespans, and if ceil is false 28, 91, 365
+        granularity (str): The granularity.
+        ceil (bool): In the case the unit is month, quarter or year. Ceil = True will use 31, 92, 366 days for these timespans, and if ceil is false 28, 91, 365
 
     Returns:
-        A timespan for the granularity
+        timedelta: A timespan for the granularity
     """
     multiplier, unit = get_granularity_multiplier_and_unit(granularity, standardize=True)
     if unit in {"w", "d", "h", "m", "s"}:

scripts/custom_checks/docstrings.py

@@ -11,14 +11,17 @@
 
 from cognite.client.data_classes import TimeSeries
 from cognite.client.data_classes.data_modeling.query import Query
+from cognite.client.testing import monkeypatch_cognite_client
 
-EXCEPTIONS = {
+FUNC_EXCEPTIONS = {}
+CLS_METHOD_EXCEPTIONS = {
     (Query, "__init__"),  # Reason: Uses a parameter 'with_'; and we need to escape the underscore
 }
 
 # Helper for testing specific class + method/property:
 TESTING = False
-ONLY_RUN = {
+ONLY_RUN_FUNCS = {monkeypatch_cognite_client}
+ONLY_RUN_CLS_METHODS = {
     (TimeSeries, "latest"),  # Just an example
 }
 
@@ -96,12 +99,12 @@ def fix_literal(s):
         annots = {}
         if isinstance(method, property):
             method_signature = inspect.signature(lambda: ...)  # just 'self' anyways
-            return_annot = method.fget.__annotations__.get("return", inspect._empty)
+            return_annot = method.fget.__annotations__.get("return", inspect.Signature.empty)
         else:
             method_signature = inspect.signature(method)
             return_annot = method_signature.return_annotation
 
-        if return_annot is inspect._empty:
+        if return_annot is inspect.Signature.empty:
             raise ValueError("Missing return type annotation")
 
         for var_name, param in method_signature.parameters.items():
@@ -279,12 +282,11 @@ def create_docstring(self):
 
         return "\n".join(final_doc_lines)
 
-    def update_py_file(self, cls, attr) -> str:
-        source_code = (path := Path(inspect.getfile(cls))).read_text()
+    def update_py_file(self, cls_or_fn, method_description) -> str:
+        source_code = (path := Path(inspect.getsourcefile(cls_or_fn))).read_text()
 
-        was_tested = f"{cls.__name__}.{attr}"
         if (n_matches := source_code.count(self.original_doc)) == 0:
-            return f"Couldn't fix docstring for '{was_tested}', as the old doc was not found in the file"
+            return f"Couldn't fix docstring for '{method_description}', as the old doc was not found in the file"
 
         elif n_matches == 1:
             new_docstr = self.create_docstring()
@@ -297,34 +299,37 @@ def update_py_file(self, cls, attr) -> str:
                 )
 
             path.write_text(source_code.replace(self.original_doc, new_docstr))
-            return f"Fixed docstring for '{was_tested}'"
+            return f"Fixed docstring for '{method_description}'"
 
         else:
-            return f"Couldn't fix docstring for '{was_tested}', as the old doc was not unique to the file"
+            return f"Couldn't fix docstring for '{method_description}', as the old doc was not unique to the file"
 
 
 def get_all_non_inherited_attributes(cls):
-    return [
-        (attr, method)
-        for attr, method in inspect.getmembers(
-            cls, predicate=lambda method: inspect.isfunction(method) or isinstance(method, property)
-        )
-        if attr in cls.__dict__
-    ]
+    def predicate(obj):
+        return inspect.isfunction(obj) or isinstance(obj, property)
+
+    return [(attr, method) for attr, method in inspect.getmembers(cls, predicate=predicate) if attr in cls.__dict__]
 
 
-def format_docstring(cls) -> list[str]:
+def format_docstring(cls_or_fn):
+    is_func = inspect.isfunction(cls_or_fn)
+    return {True: format_docstring_function, False: format_docstring_class_methods}[is_func](cls_or_fn)
+
+
+def format_docstring_class_methods(cls) -> list[str]:
     failed = []
     for attr, method in get_all_non_inherited_attributes(cls):
-        if (cls, attr) in EXCEPTIONS:
+        if (cls, attr) in CLS_METHOD_EXCEPTIONS:
             continue
 
-        if TESTING and (cls, attr) not in ONLY_RUN:
+        if TESTING and (cls, attr) not in ONLY_RUN_CLS_METHODS:
             continue
 
         # The __init__ method is documented in the class level docstring
         is_init = attr == "__init__"
         doc = cls.__doc__ if is_init else method.__doc__
+        method_description = f"{cls.__name__}.{attr}"
 
         if not doc or (is_init and is_dataclass(cls)):
             continue
@@ -335,17 +340,40 @@ def format_docstring(cls) -> list[str]:
             continue
         except (ValueError, IndexError) as e:
             failed.append(
-                f"Couldn't parse parameters in docstring for '{cls.__name__}.{attr}', "
+                f"Couldn't parse parameters in docstring for '{method_description}', "
                 f"please inspect manually. Reason: {e}"
             )
             continue
 
         if not doc_fmt.docstring_is_correct():
-            if err_msg := doc_fmt.update_py_file(cls, attr):
+            if err_msg := doc_fmt.update_py_file(cls, method_description):
                 failed.append(err_msg)
     return failed
 
 
+def format_docstring_function(fn) -> list[str]:
+    # Unwrap is needed for decorated functions to avoid getting the docstring (or later, the file to fix)
+    # of the decorator function!
+    fn = inspect.unwrap(fn)
+
+    if fn in FUNC_EXCEPTIONS or TESTING and fn not in ONLY_RUN_FUNCS or not (doc := fn.__doc__):
+        return []
+
+    fn_description = f"function: {fn.__name__}"
+
+    try:
+        doc_fmt = DocstrFormatter(doc, fn)
+    except FalsePositiveDocstring:
+        return []
+    except (ValueError, IndexError) as e:
+        return [f"Couldn't parse parameters in docstring for '{fn_description}', please inspect manually. Reason: {e}"]
+
+    if not doc_fmt.docstring_is_correct():
+        if err_msg := doc_fmt.update_py_file(fn, fn_description):
+            return [err_msg]
+    return []
+
+
 def find_all_classes_and_funcs_in_sdk():
     def predicate(obj):
         return inspect.isclass(obj) or inspect.isfunction(obj)