Remove remaining pre-3.10 typing syntax

Author: amas0

cmdstanpy/stanfit/laplace.py

@@ -2,7 +2,7 @@
 Container for the result of running a laplace approximation.
 """
 
-from typing import Any, Hashable, MutableMapping, Optional, Union
+from typing import Any, Hashable, MutableMapping
 
 import numpy as np
 import pandas as pd
@@ -41,8 +41,8 @@ def __init__(self, runset: RunSet, mode: CmdStanMLE) -> None:
         self._metadata = InferenceMetadata.from_csv(self._runset.csv_files[0])
 
     def create_inits(
-        self, seed: Optional[int] = None, chains: int = 4
-    ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]:
+        self, seed: int | None = None, chains: int = 4
+    ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
         """
         Create initial values for the parameters of the model
         by randomly selecting draws from the Laplace approximation.
@@ -168,7 +168,7 @@ def draws(self) -> np.ndarray:
 
     def draws_pd(
         self,
-        vars: Union[list[str], str, None] = None,
+        vars: list[str] | str | None = None,
     ) -> pd.DataFrame:
         if vars is not None:
             if isinstance(vars, str):
@@ -197,7 +197,7 @@ def draws_pd(
 
     def draws_xr(
         self,
-        vars: Union[str, list[str], None] = None,
+        vars: str | list[str] | None = None,
     ) -> "xr.Dataset":
         """
         Returns the sampler draws as a xarray Dataset.
@@ -308,7 +308,7 @@ def column_names(self) -> tuple[str, ...]:
         """
         return self._metadata.column_names
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Move output CSV files to specified directory.  If files were
         written to the temporary session directory, clean filename.

cmdstanpy/stanfit/mcmc.py

@@ -5,7 +5,7 @@
 import math
 import os
 from io import StringIO
-from typing import Any, Hashable, MutableMapping, Optional, Sequence, Union
+from typing import Any, Hashable, MutableMapping, Sequence
 
 import numpy as np
 import pandas as pd
@@ -97,8 +97,8 @@ def __init__(
             self._check_sampler_diagnostics()
 
     def create_inits(
-        self, seed: Optional[int] = None, chains: int = 4
-    ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]:
+        self, seed: int | None = None, chains: int = 4
+    ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
         """
         Create initial values for the parameters of the model by
         randomly selecting draws from the MCMC samples. If the samples
@@ -211,7 +211,7 @@ def column_names(self) -> tuple[str, ...]:
         return self._metadata.column_names
 
     @property
-    def metric_type(self) -> Optional[str]:
+    def metric_type(self) -> str | None:
         """
         Metric type used for adaptation, either 'diag_e' or 'dense_e', according
         to CmdStan arg 'metric'.
@@ -225,7 +225,7 @@ def metric_type(self) -> Optional[str]:
 
     # TODO(2.0): remove
     @property
-    def metric(self) -> Optional[np.ndarray]:
+    def metric(self) -> np.ndarray | None:
         """Deprecated. Use ``.inv_metric`` instead."""
         get_logger().warning(
             'The "metric" property is deprecated, use "inv_metric" instead. '
@@ -234,7 +234,7 @@ def metric(self) -> Optional[np.ndarray]:
         return self.inv_metric
 
     @property
-    def inv_metric(self) -> Optional[np.ndarray]:
+    def inv_metric(self) -> np.ndarray | None:
         """
         Inverse mass matrix used by sampler for each chain.
         Returns a ``nchains x nparams`` array when metric_type is 'diag_e',
@@ -248,7 +248,7 @@ def inv_metric(self) -> Optional[np.ndarray]:
         return self._metric
 
     @property
-    def step_size(self) -> Optional[np.ndarray]:
+    def step_size(self) -> np.ndarray | None:
         """
         Step size used by sampler for each chain.
         When sampler algorithm 'fixed_param' is specified, step size is None.
@@ -264,15 +264,15 @@ def thin(self) -> int:
         return self._thin
 
     @property
-    def divergences(self) -> Optional[np.ndarray]:
+    def divergences(self) -> np.ndarray | None:
         """
         Per-chain total number of post-warmup divergent iterations.
         When sampler algorithm 'fixed_param' is specified, returns None.
         """
         return self._divergences if not self._is_fixed_param else None
 
     @property
-    def max_treedepths(self) -> Optional[np.ndarray]:
+    def max_treedepths(self) -> np.ndarray | None:
         """
         Per-chain total number of post-warmup iterations where the NUTS sampler
         reached the maximum allowed treedepth.
@@ -564,7 +564,7 @@ def summary(
         summary_data.index.name = None
         return summary_data[mask]
 
-    def diagnose(self) -> Optional[str]:
+    def diagnose(self) -> str | None:
         """
         Run cmdstan/bin/diagnose over all output CSV files,
         return console output.
@@ -586,7 +586,7 @@ def diagnose(self) -> Optional[str]:
 
     def draws_pd(
         self,
-        vars: Union[list[str], str, None] = None,
+        vars: list[str] | str | None = None,
         inc_warmup: bool = False,
     ) -> pd.DataFrame:
         """
@@ -664,7 +664,7 @@ def draws_pd(
         )[cols]
 
     def draws_xr(
-        self, vars: Union[str, list[str], None] = None, inc_warmup: bool = False
+        self, vars: str | list[str] | None = None, inc_warmup: bool = False
     ) -> "xr.Dataset":
         """
         Returns the sampler draws as a xarray Dataset.
@@ -822,7 +822,7 @@ def method_variables(self) -> dict[str, np.ndarray]:
             for name, var in self._metadata.method_vars.items()
         }
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Move output CSV files to specified directory.  If files were
         written to the temporary session directory, clean filename.

cmdstanpy/stanfit/metadata.py

@@ -2,7 +2,7 @@
 
 import copy
 import os
-from typing import Any, Iterator, Union
+from typing import Any, Iterator
 
 import stanio
 
@@ -17,7 +17,7 @@ class InferenceMetadata:
     """
 
     def __init__(
-        self, config: dict[str, Union[str, int, float, tuple[str, ...]]]
+        self, config: dict[str, str | int | float | tuple[str, ...]]
     ) -> None:
         """Initialize object from CSV headers"""
         self._cmdstan_config = config
@@ -33,7 +33,7 @@ def __init__(
 
     @classmethod
     def from_csv(
-        cls, stan_csv: Union[str, os.PathLike, Iterator[bytes]]
+        cls, stan_csv: str | os.PathLike | Iterator[bytes]
     ) -> 'InferenceMetadata':
         try:
             comments, header, _ = stancsv.parse_comments_header_and_draws(
@@ -48,7 +48,7 @@ def from_csv(
     def __repr__(self) -> str:
         return 'Metadata:\n{}\n'.format(self._cmdstan_config)
 
-    def __getitem__(self, key: str) -> Union[str, int, float, tuple[str, ...]]:
+    def __getitem__(self, key: str) -> str | int | float | tuple[str, ...]:
         return self._cmdstan_config[key]
 
     @property

cmdstanpy/stanfit/mle.py

@@ -1,7 +1,6 @@
 """Container for the result of running optimization"""
 
 from collections import OrderedDict
-from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -58,7 +57,7 @@ def __init__(self, runset: RunSet) -> None:
             self._all_iters: np.ndarray = all_draws
 
     def create_inits(
-        self, seed: Optional[int] = None, chains: int = 4
+        self, seed: int | None = None, chains: int = 4
     ) -> dict[str, np.ndarray]:
         """
         Create initial values for the parameters of the model
@@ -134,7 +133,7 @@ def optimized_params_np(self) -> np.ndarray:
         return self._mle
 
     @property
-    def optimized_iterations_np(self) -> Optional[np.ndarray]:
+    def optimized_iterations_np(self) -> np.ndarray | None:
         """
         Returns all saved iterations from the optimizer and final estimate
         as a numpy.ndarray which contains all optimizer outputs, i.e.,
@@ -167,7 +166,7 @@ def optimized_params_pd(self) -> pd.DataFrame:
         return pd.DataFrame([self._mle], columns=self.column_names)
 
     @property
-    def optimized_iterations_pd(self) -> Optional[pd.DataFrame]:
+    def optimized_iterations_pd(self) -> pd.DataFrame | None:
         """
         Returns all saved iterations from the optimizer and final estimate
         as a pandas.DataFrame which contains all optimizer outputs, i.e.,
@@ -294,7 +293,7 @@ def stan_variables(
             )
         return result
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Move output CSV files to specified directory.  If files were
         written to the temporary session directory, clean filename.

cmdstanpy/stanfit/pathfinder.py

@@ -2,8 +2,6 @@
 Container for the result of running Pathfinder.
 """
 
-from typing import Optional, Union
-
 import numpy as np
 
 from cmdstanpy.cmdstan_args import Method
@@ -30,8 +28,8 @@ def __init__(self, runset: RunSet):
         self._metadata = InferenceMetadata.from_csv(self._runset.csv_files[0])
 
     def create_inits(
-        self, seed: Optional[int] = None, chains: int = 4
-    ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]:
+        self, seed: int | None = None, chains: int = 4
+    ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
         """
         Create initial values for the parameters of the model
         by randomly selecting draws from the Pathfinder approximation.
@@ -216,7 +214,7 @@ def is_resampled(self) -> bool:
             in (1, 'true')
         )
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Move output CSV files to specified directory.  If files were
         written to the temporary session directory, clean filename.

cmdstanpy/stanfit/runset.py

@@ -9,7 +9,6 @@
 import tempfile
 from datetime import datetime
 from time import time
-from typing import Optional
 
 from cmdstanpy import _TMPDIR
 from cmdstanpy.cmdstan_args import CmdStanArgs, Method
@@ -31,7 +30,7 @@ def __init__(
         args: CmdStanArgs,
         chains: int = 1,
         *,
-        chain_ids: Optional[list[int]] = None,
+        chain_ids: list[int] | None = None,
         time_fmt: str = "%Y%m%d%H%M%S",
         one_process_per_chain: bool = True,
     ) -> None:
@@ -162,23 +161,29 @@ def cmd(self, idx: int) -> list[str]:
             return self._args.compose_command(
                 idx,
                 csv_file=self.csv_files[idx],
-                diagnostic_file=self.diagnostic_files[idx]
-                if self._args.save_latent_dynamics
-                else None,
-                profile_file=self.profile_files[idx]
-                if self._args.save_profile
-                else None,
+                diagnostic_file=(
+                    self.diagnostic_files[idx]
+                    if self._args.save_latent_dynamics
+                    else None
+                ),
+                profile_file=(
+                    self.profile_files[idx] if self._args.save_profile else None
+                ),
             )
         else:
             return self._args.compose_command(
                 idx,
                 csv_file=self.file_path('.csv'),
-                diagnostic_file=self.file_path(".csv", extra="-diagnostic")
-                if self._args.save_latent_dynamics
-                else None,
-                profile_file=self.file_path(".csv", extra="-profile")
-                if self._args.save_profile
-                else None,
+                diagnostic_file=(
+                    self.file_path(".csv", extra="-diagnostic")
+                    if self._args.save_latent_dynamics
+                    else None
+                ),
+                profile_file=(
+                    self.file_path(".csv", extra="-profile")
+                    if self._args.save_profile
+                    else None
+                ),
             )
 
     @property
@@ -213,7 +218,7 @@ def profile_files(self) -> list[str]:
 
     # pylint: disable=invalid-name
     def file_path(
-        self, suffix: str, *, extra: str = "", id: Optional[int] = None
+        self, suffix: str, *, extra: str = "", id: int | None = None
     ) -> str:
         if id is not None:
             suffix = f"_{id}{suffix}"
@@ -256,7 +261,7 @@ def get_err_msgs(self) -> str:
                             msgs.append('\n\t'.join(errors))
         return '\n'.join(msgs)
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Moves CSV files to specified directory.
 

cmdstanpy/stanfit/vb.py

@@ -1,7 +1,6 @@
 """Container for the results of running autodiff variational inference"""
 
 from collections import OrderedDict
-from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -53,8 +52,8 @@ def __init__(self, runset: RunSet) -> None:
         self._variational_sample: np.ndarray = draws_np[1:]
 
     def create_inits(
-        self, seed: Optional[int] = None, chains: int = 4
-    ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]:
+        self, seed: int | None = None, chains: int = 4
+    ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
         """
         Create initial values for the parameters of the model
         by randomly selecting draws from the variational approximation
@@ -248,7 +247,7 @@ def variational_sample_pd(self) -> pd.DataFrame:
         """
         return pd.DataFrame(self._variational_sample, columns=self.column_names)
 
-    def save_csvfiles(self, dir: Optional[str] = None) -> None:
+    def save_csvfiles(self, dir: str | None = None) -> None:
         """
         Move output CSV files to specified directory.  If files were
         written to the temporary session directory, clean filename.

test/test_log_prob.py

@@ -4,7 +4,6 @@
 import os
 import re
 from test import check_present
-from typing import List, Optional
 
 import numpy as np
 import pytest
@@ -34,7 +33,7 @@
     ],
 )
 def test_lp_good(
-    sig_figs: Optional[int], expected: List[str], expected_unadjusted: List[str]
+    sig_figs: int | None, expected: list[str], expected_unadjusted: list[str]
 ) -> None:
     model = CmdStanModel(stan_file=BERN_STAN)
     params = {"theta": 0.34903938392023830482}