diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index af9f3841..461da83d 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -159,6 +159,19 @@ def get_study_summary(storage: BaseStorage, study_id: int) -> Optional[StudySumm return None +def create_new_study( + storage: BaseStorage, study_name: str, directions: List[StudyDirection] +) -> int: + if version.parse(optuna_ver) >= version.Version("3.1.0.dev") and version.parse( + optuna_ver + ) != version.Version("3.1.0b0"): + study_id = storage.create_new_study(study_name, directions=directions) # type: ignore + else: + study_id = storage.create_new_study(study_name) + storage.set_study_directions(study_id, directions) # type: ignore + return study_id + + def get_trials(storage: BaseStorage, study_id: int, ttl_seconds: int = 10) -> List[FrozenTrial]: with trials_cache_lock: trials = trials_cache.get(study_id, None) @@ -241,7 +254,10 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle: @json_api_view def create_study() -> BottleViewReturn: study_name = request.json.get("study_name", None) - directions = request.json.get("directions", []) + directions = [ + StudyDirection.MAXIMIZE if d.lower() == "maximize" else StudyDirection.MINIMIZE + for d in request.json.get("directions", []) + ] if ( study_name is None or len(directions) == 0 @@ -251,27 +267,7 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle: return {"reason": "You need to set study_name and direction"} try: - if version.parse(optuna_ver) >= version.Version("3.1.0.dev"): - study_id = storage.create_new_study( - study_name, - directions=[ - StudyDirection.MAXIMIZE - if d.lower() == "maximize" - else StudyDirection.MINIMIZE - for d in directions - ], - ) # type: ignore - else: - study_id = storage.create_new_study(study_name) - storage.set_study_directions( - study_id, - [ - StudyDirection.MAXIMIZE - if d.lower() == "maximize" - else StudyDirection.MINIMIZE - for d in directions - ], - ) # type: ignore + study_id = create_new_study(storage, study_name, directions) except DuplicatedStudyError: response.status = 400 # Bad request return {"reason": f"'{study_name}' already exists"} diff --git a/python_tests/test_api.py b/python_tests/test_api.py index 4f12a2e3..ceb8f60a 100644 --- a/python_tests/test_api.py +++ b/python_tests/test_api.py @@ -4,9 +4,8 @@ from unittest import TestCase import optuna from optuna import get_all_study_summaries from optuna.study import StudyDirection -from optuna.version import __version__ as optuna_ver from optuna_dashboard._app import create_app -from packaging import version +from optuna_dashboard._app import create_new_study from .wsgi_client import send_request @@ -19,17 +18,8 @@ def objective(trial: optuna.trial.Trial) -> float: class APITestCase(TestCase): def test_get_study_summaries(self) -> None: storage = optuna.storages.InMemoryStorage() - - if version.parse(optuna_ver) >= version.Version("3.1.0.dev"): - storage.create_new_study( - study_name="foo1", directions=[StudyDirection.MINIMIZE] - ) # type: ignore - storage.create_new_study( - study_name="foo2", directions=[StudyDirection.MINIMIZE] - ) # type: ignore - else: - storage.create_new_study("foo1") - storage.create_new_study("foo2") + create_new_study(storage, "foo1", [StudyDirection.MINIMIZE]) + create_new_study(storage, "foo2", [StudyDirection.MINIMIZE]) app = create_app(storage) status, _, body = send_request( @@ -138,12 +128,7 @@ class APITestCase(TestCase): def test_create_study_duplicated(self) -> None: storage = optuna.storages.InMemoryStorage() - if version.parse(optuna_ver) >= version.Version("3.1.0.dev"): - storage.create_new_study( - study_name="foo", directions=[StudyDirection.MINIMIZE] - ) # type: ignore - else: - storage.create_new_study("foo") + create_new_study(storage, "foo", [StudyDirection.MINIMIZE]) self.assertEqual(len(get_all_study_summaries(storage)), 1) app = create_app(storage) @@ -163,16 +148,8 @@ class APITestCase(TestCase): def test_delete_study(self) -> None: storage = optuna.storages.InMemoryStorage() - if version.parse(optuna_ver) >= version.Version("3.1.0.dev"): - storage.create_new_study( - study_name="foo1", directions=[StudyDirection.MINIMIZE] - ) # type: ignore - storage.create_new_study( - study_name="foo2", directions=[StudyDirection.MINIMIZE] - ) # type: ignore - else: - storage.create_new_study(study_name="foo1") - storage.create_new_study(study_name="foo2") + create_new_study(storage, "foo1", [StudyDirection.MINIMIZE]) + create_new_study(storage, "foo2", [StudyDirection.MINIMIZE]) self.assertEqual(len(get_all_study_summaries(storage)), 2) app = create_app(storage)