Merge branch 'main' into fix_marginal_cost

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]

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:

assume/cli.py → cli.py

@@ -1,74 +1,122 @@
+#!/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",
         "--scenario",
         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(" "))

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
 =========================================
 

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"]

tests/test_integration_cli.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from assume.cli import cli
+from cli import cli
 
 
 @pytest.mark.slow