From c32a623a5e556b4fa11bb84ee18882e6b9fcb34c Mon Sep 17 00:00:00 2001 From: Florian Maurer Date: Mon, 13 Nov 2023 12:01:15 +0100 Subject: [PATCH] add fancy completion using argcomplete (#251) This suggests possible database paths and scenarios, as well as parameters for case_study based on the context. By moving the script to the top folder, assume does not need to be imported if only calling the --help for example. This should also work on macOS by running: `pip install -e .` `eval "$(register-python-argcomplete assume)"` And also for Powershell on Windows using: `register-python-argcomplete --shell powershell assume | Out-String | Invoke-Expression` (though I did not test it) --- assume/common/scenario_loader.py | 3 +- assume/world.py | 4 +- assume/cli.py => cli.py | 74 ++++++++++++++++++++++++++------ docs/source/installation.rst | 16 +++++++ pyproject.toml | 7 +-- tests/test_integration_cli.py | 2 +- 6 files changed, 87 insertions(+), 19 deletions(-) rename assume/cli.py => cli.py (52%) diff --git a/assume/common/scenario_loader.py b/assume/common/scenario_loader.py index ba1c383f..e149b10c 100644 --- a/assume/common/scenario_loader.py +++ b/assume/common/scenario_loader.py @@ -254,7 +254,8 @@ async def load_scenario_folder_async( # load the config file path = f"{inputs_path}/{scenario}" - config = yaml.safe_load(open(f"{path}/config.yaml", "r")) + with open(f"{path}/config.yaml", "r") as f: + config = yaml.safe_load(f) if not study_case: study_case = list(config.keys())[0] config = config[study_case] diff --git a/assume/world.py b/assume/world.py index 79158265..a91f86aa 100644 --- a/assume/world.py +++ b/assume/world.py @@ -17,7 +17,7 @@ from mango.util.clock import ExternalClock from mango.util.distributed_clock import DistributedClockAgent, DistributedClockManager from mango.util.termination_detection import tasks_complete_or_sleeping -from sqlalchemy import create_engine +from sqlalchemy import create_engine, make_url from sqlalchemy.exc import OperationalError from tqdm import tqdm @@ -59,7 +59,7 @@ def __init__( self.export_csv_path = export_csv_path # intialize db connection at beginning of simulation if database_uri: - self.db = create_engine(database_uri) + self.db = create_engine(make_url(database_uri)) connected = False while not connected: try: diff --git a/assume/cli.py b/cli.py similarity index 52% rename from assume/cli.py rename to cli.py index b9799a80..1d492b62 100644 --- a/assume/cli.py +++ b/cli.py @@ -1,22 +1,65 @@ +#!/usr/bin/env python +# PYTHON_ARGCOMPLETE_OK + # SPDX-FileCopyrightText: ASSUME Developers # # SPDX-License-Identifier: AGPL-3.0-or-later import argparse +import logging import os import sys +from pathlib import Path -from assume import World, load_scenario_folder, run_learning +import argcomplete +import yaml +from sqlalchemy import make_url os.makedirs("./examples/outputs", exist_ok=True) os.makedirs("./examples/local_db", exist_ok=True) +def db_uri_completer(prefix, parsed_args, **kwargs): + return { + "sqlite://example.db": "example", + f"sqlite://examples/local_db/{parsed_args.scenario}.db": "current scenario", + "sqlite:///": "in-memory", + "postgresql://assume:assume@localhost:5432/assume": "localhost", + "postgresql://assume:assume@assume_db:5432/assume": "docker", + "mysql://username:password@localhost:3306/database": "mysql", + } + + +def config_directory_completer(prefix, parsed_args, **kwargs): + directory = Path(parsed_args.input_path) + if directory.is_dir(): + config_folders = [ + folder + for folder in directory.iterdir() + if folder.is_dir() and (folder / "config.yaml").exists() + ] + return [ + folder.name for folder in config_folders if folder.name.startswith(prefix) + ] + return [""] + + +def config_case_completer(prefix, parsed_args, **kwargs): + config_file = ( + Path(parsed_args.input_path) / Path(parsed_args.scenario) / "config.yaml" + ) + if config_file.is_file(): + with open(str(config_file), "r") as f: + config = yaml.safe_load(f) + return list(config.keys()) + return [""] + + def cli(args=None): if not args: args = sys.argv[1:] parser = argparse.ArgumentParser( - description="Command Line Interface for ASSUME simulations" + description="Command Line Interface for ASSUME simulations", ) parser.add_argument( "-s", @@ -24,51 +67,56 @@ def cli(args=None): help="name of the scenario file which should be used", default="example_01a", type=str, - ) + ).completer = config_directory_completer parser.add_argument( "-c", "--case-study", help="name of the case in that scenario which should be simulated", default="", type=str, - ) + ).completer = config_case_completer parser.add_argument( "-csv", "--csv-export-path", help="optional path to the csv export", default="", type=str, - ) + ).completer = argcomplete.DirectoriesCompleter() parser.add_argument( "-db", "--db-uri", help="uri string for a database", default="", type=str, - ) + ).completer = db_uri_completer parser.add_argument( - "-input", + "-i", "--input-path", help="path to the input folder", default="examples/inputs", type=str, - ) + ).completer = argcomplete.DirectoriesCompleter() parser.add_argument( "-l", "--loglevel", help="logging level used for file log", default="INFO", type=str, + metavar="LOGLEVEL", + choices=set(logging._nameToLevel.keys()), ) + argcomplete.autocomplete(parser) args = parser.parse_args(args) name = args.scenario if args.db_uri: - db_uri = args.db_uri + db_uri = make_url(args.db_uri) else: db_uri = f"sqlite:///./examples/local_db/{name}.db" try: + from assume import World, load_scenario_folder, run_learning + world = World( database_uri=db_uri, export_csv_path=args.csv_export_path, @@ -96,6 +144,8 @@ def cli(args=None): if __name__ == "__main__": - # cli() - args = "-s example_01_rl -db postgresql://assume:assume@localhost:5432/assume" - cli(args.split(" ")) + cli() + + # args = "-s example_02 -db postgresql://assume:assume@localhost:5432/assume" + + # cli(args.split(" ")) diff --git a/docs/source/installation.rst b/docs/source/installation.rst index e1b41fc5..3449a587 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -75,6 +75,22 @@ To install with testing capabilities:: pip install assume-framework[test] +Install Tab-Completion +----------------- + +ASSUME uses `argcomplete` for argument completion on the CLI. + +On Windows, one needs to run: + +`register-python-argcomplete --shell powershell assume | Out-String | Invoke-Expression` + +in the used conda environment, to install tab completions. + +On Bash or zsh (Linux and Mac) run the following in the correct conda environment with assume and argcomplete installed: + +`eval "$(register-python-argcomplete assume)"` + + Install using Docker ========================================= diff --git a/pyproject.toml b/pyproject.toml index c7ef06c5..7514104f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,8 @@ packages = [ [tool.poetry.dependencies] python = "^3.10" +argcomplete = "^3.1.4" +nest-asyncio = "^1.5.6" paho-mqtt = "^1.5.1" mango-agents-assume = "^1.1.1-1" tqdm = "^4.64.1" @@ -39,7 +41,7 @@ pandas = {version = "^2.0.0"} psycopg2-binary = "^2.9.5" pyomo = "^6.6.1" pyyaml = "^6.0" -nest-asyncio = "^1.5.6" +pyyaml-include = "^1.3.1" black = {version = "^23.3.0", optional = true} isort = {version = "^5.12.0", optional = true} mypy = {version = "^1.1.1", optional = true} @@ -48,7 +50,6 @@ pytest = {version = "^7.2.2", optional = true} pytest-cov = {version = "^4.1.0", optional = true} pytest-asyncio = {version = "^0.21.1", optional = true} torch = {version = "^2.0.1", optional = true} - glpk = {version = "^0.4.7", optional = true} [tool.poetry.group.dev.dependencies] @@ -67,7 +68,7 @@ test = ["black", "isort", "matplotlib", "pytest", "pytest-cov", "pytest-asyncio" [tool.poetry.scripts] -assume = "assume.cli:cli" +assume = "cli:cli" [build-system] requires = ["poetry-core"] diff --git a/tests/test_integration_cli.py b/tests/test_integration_cli.py index 994a36ee..036d4fb7 100644 --- a/tests/test_integration_cli.py +++ b/tests/test_integration_cli.py @@ -4,7 +4,7 @@ import pytest -from assume.cli import cli +from cli import cli @pytest.mark.slow