From 7a05d63b232a751311364e79648581322f3cda60 Mon Sep 17 00:00:00 2001 From: Eero Vaher Date: Mon, 11 May 2026 19:46:40 +0200 Subject: [PATCH] Replace `TestFunction.var_dtypes` The attribute was initialized as an empty `dict` in `TestFunction.__init__()` and filled out as a side-effect of calling `TestFunction.define_arrays()`, but in practice its job can be done by a `frozenset` that can be easily constructed in `TestFunction.__init__()`. --- erfa_generator.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/erfa_generator.py b/erfa_generator.py index d6718e3..e2272bd 100644 --- a/erfa_generator.py +++ b/erfa_generator.py @@ -360,8 +360,9 @@ def __init__(self, name, t_erfa_c, nin, ninout, nout): self.nin = nin self.ninout = ninout self.nout = nout - # Dict of dtypes for variables, filled by define_arrays(). - self.var_dtypes = {} + self.dt_pv_vars: Final = frozenset( + re.findall(r"(\w+)\[2\]\[3\]", search.group(1)) + ) @classmethod def from_function(cls, func, t_erfa_c): @@ -429,7 +430,6 @@ def define_arrays(self, line): v_dtype = v.dtype v_shape = v.shape if v.signature_shape != '()' else '()' extra = "" - self.var_dtypes[name] = v_dtype v_dtype = "float" if v_dtype == "dt_double" else "erfa_ufunc." + v_dtype defines.append(f"{name} = np.empty({v_shape}, {v_dtype}){extra}") @@ -540,8 +540,7 @@ def to_python(self): # that were not caught by the general replacement above (e.g., # with names not equal to 'pv') name, _, rest = line.partition('[') - if (rest and rest[0] in '01' and name in self.var_dtypes - and self.var_dtypes[name] == 'dt_pv'): + if name in self.dt_pv_vars and rest.startswith(("0", "1")): line = name + "[" + ("'p'" if rest[0] == "0" else "'v'") + rest[1:] out.append(line)