diff --git a/.gitignore b/.gitignore index 348f767..cb7ba4a 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,6 @@ build/ .coverage .tox *.log +.history/* +.aiida_projects/ +.vscode/ diff --git a/aiida_project/commands/main.py b/aiida_project/commands/main.py index b174589..d061314 100644 --- a/aiida_project/commands/main.py +++ b/aiida_project/commands/main.py @@ -7,24 +7,9 @@ from rich import print, prompt from typing_extensions import Annotated -from ..config import ShellType +from ..config import ShellGenerator, ShellType from ..project import EngineType, load_project_class -CDA_FUNCTION = """ -cda () { - source $aiida_venv_dir/$1/bin/activate - cd $aiida_project_dir/$1 -} -""" - -ACTIVATE_AIIDA_SH = """ -export AIIDA_PATH={path} -eval "$(_VERDI_COMPLETE={shell}_source verdi)" -""" - -DEACTIVATE_AIIDA_SH = "unset AIIDA_PATH" - - app = typer.Typer(pretty_exceptions_show_locals=False) @@ -40,7 +25,12 @@ def init(shell: Optional[ShellType] = None): """Initialisation of the `aiida-project` setup.""" from ..config import ProjectConfig - config = ProjectConfig() + # ! When instantiated without arguments, default values aren't picked up for some reason... + # config = ProjectConfig() + config = ProjectConfig( + aiida_venv_dir = Path(Path.home(), ".aiida_venvs"), + aiida_project_dir = Path(Path.home(), "aiida_projects") + ) shell_str = shell.value if shell else None @@ -52,21 +42,25 @@ def init(shell: Optional[ShellType] = None): prompt=prompt_message, choices=[shell_type.value for shell_type in ShellType] ) - if shell_str == "fish": - print( - "[bold red]Error:[/] `fish` is not yet supported. " - "If this is Julian: you better get to it 😜" # Nhehehe - ) - return + shellgenerator = ShellGenerator(shell_str=shell_str) - config.set_key("aiida_project_shell", shell_str) + config.write_key("aiida_project_shell", shell_str) - env_file_path = Path.home() / Path(f".{shell_str}rc") + # ? Add this to ShellGenerator? + if shell_str == 'fish': + env_file_path = Path.home() / Path( + os.path.join('.config', 'fish', 'conf.d', 'aiida_project.fish') + ) + if not env_file_path.exists(): + env_file_path.touch() + + else: + env_file_path = Path.home() / Path(f".{shell_str}rc") if "Created by `aiida-project init`" in env_file_path.read_text(): print( "[bold blue]Report:[/] There is already an `aiida-project` initialization comment in " - f"{env_file_path}." + f"{env_file_path}" # ? Removed trailing dot ) add_init_lines = prompt.Confirm.ask("Do you want still want to add the init lines?") else: @@ -78,15 +72,25 @@ def init(shell: Optional[ShellType] = None): f"\n# Created by `aiida-project init` on " f"{datetime.now().strftime('%d/%m/%y %H:%M')}\n" ) - handle.write(f"export $(grep -v '^#' {config.Config.env_file} | xargs)") - handle.write(CDA_FUNCTION) + # ? Pass as function argument, rather than via .format on the returned str? + handle.write( + shellgenerator.variable_export().format( + env_file_path=config.Config.env_file + ) + ) + handle.write(shellgenerator.cd_aiida()) - config.set_key( + config.write_key( + "aiida_venv_dir", + os.environ.get("WORKON_HOME", config.aiida_venv_dir.as_posix()), + ) + config.write_key("aiida_project_dir", config.aiida_project_dir.as_posix()) + config.write_key( "aiida_venv_dir", os.environ.get("WORKON_HOME", config.aiida_venv_dir.as_posix()), ) - config.set_key("aiida_project_dir", config.aiida_project_dir.as_posix()) print("✨🚀 AiiDA-project has been initialised! 🚀✨") + print("Restart your shell to let the changes take effect.") @app.command() @@ -114,9 +118,17 @@ def create( config = ProjectConfig() if config.is_not_initialised(): return + else: + config = config.from_env_file() + # raise SystemExit("Ciao 👋") venv_path = config.aiida_venv_dir / Path(name) project_path = config.aiida_project_dir / Path(name) + shell_str = config.aiida_project_shell + # ? shell_str should be set as instance variable and is used for the command + # ? selection. However, it's value is also passed to the string format method. + # ? Right now, it seems a bit clunky and duplicated... + shellgenerator = ShellGenerator(shell_str=shell_str) # Temporarily block `conda` engines until we provide support again if engine is EngineType.conda: @@ -138,9 +150,12 @@ def create( typer.echo("🔧 Adding the AiiDA environment variables to the activate script.") project.append_activate_text( - ACTIVATE_AIIDA_SH.format(path=project_path, shell=config.aiida_project_shell) + shellgenerator.activate_aiida().format( + env_file_path=project_path, + shell_str=shell_str + ) ) - project.append_deactivate_text(DEACTIVATE_AIIDA_SH) + project.append_deactivate_text(shellgenerator.deactivate_aiida()) project_dict = ProjectDict() project_dict.add_project(project) @@ -188,4 +203,4 @@ def destroy( project.destroy() project_dict.remove_project(name) - print(f"[bold green]Succes:[/bold green] Project with name {name} has been destroyed.") + print(f"[bold green]Succes:[/bold green] Project with name {name} has been destroyed.") \ No newline at end of file diff --git a/aiida_project/config.py b/aiida_project/config.py index 0280cc6..28e058f 100644 --- a/aiida_project/config.py +++ b/aiida_project/config.py @@ -1,3 +1,4 @@ +import textwrap from enum import Enum from pathlib import Path from typing import Dict, Optional, Union @@ -25,11 +26,95 @@ class ShellType(str, Enum): fish = "fish" +class ShellGenerator: + + """Class that sets shell environment variables, relevant directories, config files, etc.""" + + CD_AIIDA_BASH = """ + cda () { + source "$aiida_venv_dir/$1/bin/activate" + cd "$aiida_project_dir/$1" + } + """ + CD_AIIDA_FISH = """ + function cda + source "$aiida_venv_dir/$argv[1]/bin/activate.fish" + cd "$aiida_project_dir/$argv[1]" + end + funcsave -q cda + """ + + ACTIVATE_AIIDA_BASH = """ + export AIIDA_PATH={env_file_path} + eval "$(_VERDI_COMPLETE={shell_str}_source verdi)" + """ + + ACTIVATE_AIIDA_FISH = """ + set -Ux AIIDA_PATH {env_file_path} + eval "$(_VERDI_COMPLETE={shell_str}_source verdi)" + """ + + VARIABLE_EXPORT_BASH = "export $(grep -v '^#' {env_file_path} | xargs)" + VARIABLE_EXPORT_FISH = """ + grep -v '^#' {env_file_path} | sed "s|'||g; s|\\"||g; s|=| |" | while read -l key value + set -Ux "$key" "$value" + end + """ + + DEACTIVATE_AIIDA_BASH = "unset AIIDA_PATH" + DEACTIVATE_AIIDA_FISH = "set -e AIIDA_PATH" + + # def __init__(self, shell_str: str, env_file_path: Optional[Path] = None) -> None: + def __init__(self, shell_str: str): + self.shell_str = shell_str + # self.env_file_path = env_file_path + + def _get_script(self, bash_script: str, fish_script: str) -> str: + + """Helper function to return respective shell script based type. Avoids code duplication. + + Args: + bash_script (str): Bash codes defined above as class variables. + fish_script (str): Fish codes defined above as class variables. + + Raises: + ValueError: If shell type is not one of the supported types. + + Returns: + str: Respective shell code. + """ + + if self.shell_str in ['bash', 'zsh']: + return_script = bash_script + elif self.shell_str == 'fish': + return_script = fish_script + else: + raise ValueError(f'Invalid shell type: {self.shell_str}') + # ? Remove indentation from multiline string + return textwrap.dedent(return_script) + + def cd_aiida(self): + """Define function to `cd` into a given AiiDA project.""" + return self._get_script(self.CD_AIIDA_BASH, self.CD_AIIDA_FISH) + + def activate_aiida(self): + """Define function to activate a given AiiDA project.""" + return self._get_script(self.ACTIVATE_AIIDA_BASH, self.ACTIVATE_AIIDA_FISH) + + def variable_export(self): + """Set relevant environment variables.""" + return self._get_script(self.VARIABLE_EXPORT_BASH, self.VARIABLE_EXPORT_FISH) + + def deactivate_aiida(self): + """Unset AiiDA PATH.""" + return self._get_script(self.DEACTIVATE_AIIDA_BASH, self.DEACTIVATE_AIIDA_FISH) + + class ProjectConfig(BaseSettings): """Configuration class for configuring `aiida-project`.""" - aiida_venv_dir: Path = Path(Path.home(), ".aiida_venvs") - aiida_project_dir: Path = Path(Path.home(), "project") + aiida_venv_dir: Path = Path.home() / Path(".aiida_venvs") + aiida_project_dir: Path = Path.home() / Path("aiida_projects") # ? Make hidden? aiida_default_python_path: Optional[Path] = None aiida_project_structure: dict = DEFAULT_PROJECT_STRUCTURE aiida_project_shell: str = "bash" @@ -44,10 +129,24 @@ def is_not_initialised(self): print("[bold blue]Info:[/bold blue] Please run `aiida-project init` to get started.") return True - def set_key(self, key, value): + @classmethod + def from_env_file(cls, env_path: Optional[Path] = None): + + """Populate config instance from env file.""" + + if env_path is None: + env_path = Path.home() / Path(".aiida_project.env") + # ? Currently works only with default path of Config class... + config_dict = {k:v for k,v in dotenv.dotenv_values(env_path).items()} + config_instance = cls.parse_obj(config_dict) + return config_instance + + # ? Renamed these methods, as they actually write to the env file, and we + # ? might implement a method that sets the key of the ProjectConfig class? + def write_key(self, key, value): dotenv.set_key(self.Config.env_file, key, value) - def get_key(self, key): + def read_key(self, key): return dotenv.get_key(key)