From a0e5c96cee8b320575a292dca90e7adb3beb96c1 Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Wed, 22 Sep 2021 17:28:04 -0500 Subject: [PATCH] Add type to derivative_coeff_dict --- sumpy/tools.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/sumpy/tools.py b/sumpy/tools.py index bec9ea9b..89192f78 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -46,6 +46,7 @@ import numpy as np import sumpy.symbolic as sym import loopy as lp +from typing import Dict, Tuple, Any import logging logger = logging.getLogger(__name__) @@ -394,6 +395,9 @@ class HelmholtzDerivativeTaker(RadialDerivativeTaker): return expr +DerivativeCoeffDict = Dict[Tuple[int], Any] + + @tag_dataclass class DifferentiatedExprDerivativeTaker: """Implements the :class:`ExprDerivativeTaker` interface @@ -411,7 +415,7 @@ class DifferentiatedExprDerivativeTaker: base expression. """ taker: ExprDerivativeTaker - derivative_coeff_dict: dict + derivative_coeff_dict: DerivativeCoeffDict def diff(self, mi, save_intermediate=lambda x: x): # By passing `rscale` to the derivative taker we are taking a scaled @@ -432,7 +436,8 @@ class DifferentiatedExprDerivativeTaker: return result * save_intermediate(1 / self.taker.rscale ** max_order) -def diff_derivative_coeff_dict(derivative_coeff_dict, variable_idx, variables): +def diff_derivative_coeff_dict(derivative_coeff_dict: DerivativeCoeffDict, + variable_idx, variables): """Differentiate a derivative transformation dictionary given by *derivative_coeff_dict* using the variable given by **variable_idx** and return a new derivative transformation dictionary. -- GitLab