From d24c6d1ea331b63d92a98552193b590f5fb11627 Mon Sep 17 00:00:00 2001 From: Makoto Hiramatsu Date: Sun, 24 Apr 2022 16:21:38 +0900 Subject: [PATCH] Add annotation for distributions --- python_tests/test_search_space.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python_tests/test_search_space.py b/python_tests/test_search_space.py index df7fe9e2..44b460ea 100644 --- a/python_tests/test_search_space.py +++ b/python_tests/test_search_space.py @@ -1,9 +1,12 @@ +from typing import Dict +from typing import List from unittest import TestCase import warnings import optuna from optuna import create_trial from optuna.distributions import UniformDistribution +from optuna.distributions import BaseDistribution from optuna.exceptions import ExperimentalWarning from optuna.trial import TrialState @@ -16,7 +19,7 @@ class SearchSpaceTestCase(TestCase): warnings.simplefilter("ignore", category=ExperimentalWarning) def test_same_distributions(self) -> None: - distributions = [ + distributions: List[Dict[str, BaseDistribution]] = [ { "x0": UniformDistribution(low=0, high=10), "x1": UniformDistribution(low=0, high=10), @@ -47,7 +50,7 @@ class SearchSpaceTestCase(TestCase): self.assertEqual(len(search_space.union), 2) def test_different_distributions(self) -> None: - distributions = [ + distributions: List[Dict[str, BaseDistribution]] = [ { "x0": UniformDistribution(low=0, high=10), "x1": UniformDistribution(low=0, high=10), @@ -78,7 +81,7 @@ class SearchSpaceTestCase(TestCase): self.assertEqual(len(search_space.union), 3) def test_dynamic_search_space(self) -> None: - distributions = [ + distributions: List[Dict[str, BaseDistribution]] = [ { "x0": UniformDistribution(low=0, high=10), "x1": UniformDistribution(low=0, high=10),