Test that numpy and other functions work in eval.
()
| 151 | |
| 152 | |
| 153 | def test_eval_functions() -> None: |
| 154 | """Test that numpy and other functions work in eval.""" |
| 155 | ds = Dataset({"a": ("x", [0.0, 1.0, 4.0])}) |
| 156 | |
| 157 | # numpy functions via np namespace should work |
| 158 | result = ds.eval("np.sqrt(a)") |
| 159 | assert_equal(result, np.sqrt(ds["a"])) |
| 160 | |
| 161 | result = ds.eval("np.sin(a) + np.cos(a)") |
| 162 | assert_equal(result, np.sin(ds["a"]) + np.cos(ds["a"])) |
| 163 | |
| 164 | # pandas namespace should work |
| 165 | result = ds.eval("pd.isna(a)") |
| 166 | # pd.isna returns ndarray, not DataArray |
| 167 | np.testing.assert_array_equal(result, pd.isna(ds["a"].values)) |
| 168 | |
| 169 | # xarray namespace should work |
| 170 | result = ds.eval("xr.where(a > 1, a, 0)") |
| 171 | |
| 172 | assert_equal(result, xr.where(ds["a"] > 1, ds["a"], 0)) |
| 173 | |
| 174 | # Common builtins should work |
| 175 | result = ds.eval("abs(a - 2)") |
| 176 | assert_equal(result, abs(ds["a"] - 2)) |
| 177 | |
| 178 | result = ds.eval("round(float(a.mean()))") |
| 179 | assert result == round(float(ds["a"].mean())) |
| 180 | |
| 181 | result = ds.eval("len(a)") |
| 182 | assert result == 3 |
| 183 | |
| 184 | result = ds.eval("pow(a, 2)") |
| 185 | assert_equal(result, ds["a"] ** 2) |
| 186 | |
| 187 | # Attribute access on DataArrays should work |
| 188 | result = ds.eval("a.values") |
| 189 | assert isinstance(result, np.ndarray) |
| 190 | |
| 191 | # Method calls on DataArrays should work |
| 192 | result = ds.eval("a.mean()") |
| 193 | assert float(result) == np.mean([0.0, 1.0, 4.0]) |
| 194 | |
| 195 | |
| 196 | def test_eval_extended_builtins() -> None: |