From 275b6cd9c00e8e126aeee1b24842d85601083a7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Sch=C3=A4fer?= <david.schaefer@ufz.de>
Date: Tue, 19 Jan 2021 16:19:30 +0100
Subject: [PATCH] residues: type hints

---
 requirements.txt       |  1 +
 saqc/funcs/residues.py | 32 +++++++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 7f0858637..101eda485 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -32,4 +32,5 @@ wcwidth==0.2.5
 scipy==1.6.0
 scikit-learn==0.23.2
 six==1.15.0
+typing_extensions==3.7.4.3
 astor==0.8.1
diff --git a/saqc/funcs/residues.py b/saqc/funcs/residues.py
index 864c5c955..7dee7eeea 100644
--- a/saqc/funcs/residues.py
+++ b/saqc/funcs/residues.py
@@ -1,15 +1,27 @@
 #! /usr/bin/env python
 # -*- coding: utf-8 -*-
+
+from typing import Tuple, Union, Optional, Callable
+from typing_extensions import Literal
+
 import numpy as np
 
+from dios import DictOfSeries
+
 from saqc.core.register import register
+from saqc.flagger.baseflagger import BaseFlagger
 from saqc.funcs.rolling import roll
 from saqc.funcs.curvefit import fitPolynomial
 
 
 @register(masking='field')
-def calculatePolynomialResidues(data, field, flagger, winsz, polydeg, numba="auto", eval_flags=True, min_periods=0,
-                                **kwargs):
+def calculatePolynomialResidues(
+        data: DictOfSeries, field: str, flagger: BaseFlagger,
+        winsz: Union[str, int], polydeg: int,
+        numba: Literal[True, False, "auto"]="auto",
+        eval_flags: bool=True,
+        min_periods: Optional[int]=0,
+        **kwargs) -> Tuple[DictOfSeries, BaseFlagger]:
     """
     Function fits a polynomial model to the data and returns the residues.
 
@@ -67,7 +79,7 @@ def calculatePolynomialResidues(data, field, flagger, winsz, polydeg, numba="aut
     eval_flags : bool, default True
         Wheather or not to assign new flags to the calculated residuals. If True, a residual gets assigned the worst
         flag present in the interval, the data for its calculation was obtained from.
-    min_periods : {int, np.nan}, default 0
+    min_periods : {int, None}, default 0
         The minimum number of periods, that has to be available in every values fitting surrounding for the polynomial
         fit to be performed. If there are not enough values, np.nan gets assigned. Default (0) results in fitting
         regardless of the number of values present (results in overfitting for too sparse intervals). To automatically
@@ -84,16 +96,22 @@ def calculatePolynomialResidues(data, field, flagger, winsz, polydeg, numba="aut
 
     """
     data, flagger = fitPolynomial(data, field, flagger, winsz, polydeg, numba=numba, eval_flags=eval_flags,
-                                  min_periods=min_periods, _return_residues=True, **kwargs)
+                                  min_periods=min_periods, return_residues=True, **kwargs)
 
     return data, flagger
 
 
 @register(masking='field')
-def calculateRollingResidues(data, field, flagger, winsz, func=np.mean, eval_flags=True, min_periods=0, center=True,
-                             **kwargs):
+def calculateRollingResidues(
+        data: DictOfSeries, field: str, flagger: BaseFlagger,
+        winsz: Union[str, int],
+        func: Callable[[np.array], np.array]=np.mean,
+        eval_flags: bool=True,
+        min_periods: Optional[int]=0,
+        center: bool=True,
+        **kwargs) -> Tuple[DictOfSeries, BaseFlagger]:
 
     data, flagger = roll(data, field, flagger, winsz, func=func, eval_flags=eval_flags,
                          min_periods=min_periods, center=center, _return_residues=True, **kwargs)
 
-    return data, flagger
\ No newline at end of file
+    return data, flagger
-- 
GitLab