diff --git a/empress/_parameter_descriptions.py b/empress/_parameter_descriptions.py index 36905f86..5df7ca25 100644 --- a/empress/_parameter_descriptions.py +++ b/empress/_parameter_descriptions.py @@ -93,3 +93,5 @@ OUTPUT_DIR = ( 'Directory to create in which an EMPress visualization will be written.' ) + +OVERWRITE = 'Determines whether to overwrite existing visualization files.' diff --git a/empress/_plot_utils.py b/empress/_plot_utils.py index 57ed8c5d..41410aac 100644 --- a/empress/_plot_utils.py +++ b/empress/_plot_utils.py @@ -39,7 +39,7 @@ def get_bp(newickfmt): return parse_newick(treefile.readline()) -def save_viz(viz, output_dir, q2=True): +def save_viz(viz, output_dir, q2=True, overwrite=False): """Saves an Empress visualization to a filepath. Parameters @@ -47,11 +47,12 @@ def save_viz(viz, output_dir, q2=True): viz : empress.Empress output_dir : str q2 : bool + overwrite : bool """ with open(os.path.join(output_dir, 'empress.html'), 'w') as htmlfile: htmlfile.write(str(viz)) - viz.copy_support_files(output_dir) + viz.copy_support_files(output_dir, overwrite=overwrite) if q2: import q2templates @@ -83,7 +84,8 @@ def prepare_pcoa(pcoa, number_of_features): return pcoa -def check_and_process_files(output_dir, tree_file, feature_metadata): +def check_and_process_files(output_dir, tree_file, feature_metadata, + overwrite=False): """Initial checks and processing of files for standalone CLI plotting. Parameters @@ -91,6 +93,7 @@ def check_and_process_files(output_dir, tree_file, feature_metadata): output_dir : str tree_file : str fm_file : str + overwrite: bool Returns ------- @@ -98,7 +101,8 @@ def check_and_process_files(output_dir, tree_file, feature_metadata): pd.DataFrame """ if os.path.isdir(output_dir): - raise OSError("Output directory already exists!") + if not overwrite: + raise OSError("Output directory already exists!") with open(str(tree_file), "r") as f: tree_newick = parse_newick(f.readline()) if feature_metadata is not None: diff --git a/empress/core.py b/empress/core.py index a86c1bd2..d3195e71 100644 --- a/empress/core.py +++ b/empress/core.py @@ -244,7 +244,7 @@ def _validate_and_match_data(self, ignore_missing_samples, ) = match_tree_and_feature_metadata(self.tree, self.features) validate_tree(self.tree) - def copy_support_files(self, target=None): + def copy_support_files(self, target=None, overwrite=False): """Copies the support files to a target directory If an ordination is included Emperor's support files will also be @@ -260,7 +260,8 @@ def copy_support_files(self, target=None): target = self.base_url # copy the required resources - copytree(SUPPORT_FILES, os.path.join(target, 'support_files')) + copytree(SUPPORT_FILES, os.path.join(target, 'support_files'), + dirs_exist_ok=overwrite) if self._emperor is not None: self._emperor.copy_support_files(os.path.join(target, diff --git a/empress/scripts/_cli.py b/empress/scripts/_cli.py index 24976cb8..2df78452 100644 --- a/empress/scripts/_cli.py +++ b/empress/scripts/_cli.py @@ -30,22 +30,26 @@ def empress(): help=desc.FM_DESC) @click.option("--shear-to-feature-metadata", required=False, default=False, help=desc.SHEAR_TO_FM, is_flag=True) +@click.option("--overwrite", required=False, default=False, + help=desc.OVERWRITE, is_flag=True) def tree_plot( tree: str, output_dir: str, feature_metadata: str, shear_to_feature_metadata: bool, + overwrite: bool ) -> None: tree_newick, fm = check_and_process_files( output_dir, tree, - feature_metadata + feature_metadata, + overwrite ) viz = Empress(tree_newick, feature_metadata=fm, shear_to_feature_metadata=shear_to_feature_metadata) - os.makedirs(output_dir) - save_viz(viz, output_dir, q2=False) + os.makedirs(output_dir, exist_ok=overwrite) + save_viz(viz, output_dir, q2=False, overwrite=overwrite) @empress.command( @@ -68,6 +72,8 @@ def tree_plot( help=desc.NUM_FEAT) @click.option("--shear-to-table", required=False, default=True, help=desc.SHEAR_TO_TBL, is_flag=True) +@click.option("--overwrite", required=False, default=False, + help=desc.OVERWRITE, is_flag=True) def community_plot( tree: str, table: str, @@ -80,11 +86,13 @@ def community_plot( filter_missing_features: bool, number_of_pcoa_features: int, shear_to_table: bool, + overwrite: bool ) -> None: tree_newick, fm = check_and_process_files( output_dir, tree, - feature_metadata + feature_metadata, + overwrite ) table = load_table(table) sample_metadata = pd.read_csv(sample_metadata, sep="\t", index_col=0) @@ -104,8 +112,8 @@ def community_plot( filter_missing_features=filter_missing_features, shear_to_table=shear_to_table, ) - os.makedirs(output_dir) - save_viz(viz, output_dir, q2=False) + os.makedirs(output_dir, exist_ok=overwrite) + save_viz(viz, output_dir, q2=False, overwrite=overwrite) if __name__ == "__main__": diff --git a/tests/python/test_cli.py b/tests/python/test_cli.py index 99919ca1..5730e08b 100644 --- a/tests/python/test_cli.py +++ b/tests/python/test_cli.py @@ -105,6 +105,18 @@ def test_existing_directory(cls): assert not os.path.isdir(f"{output_dir}/support_files") assert "empress.html" not in os.listdir(output_dir) + def test_existing_directory_overwrite(cls): + output_dir = "existing_dir_2" + os.mkdir(output_dir) + result = cls.runner.invoke( + empress, + ["community-plot", "--tree", cls.tree_loc, "--table", + cls.table_loc, "--sample-metadata", cls.sm_loc, + "--output-dir", output_dir, "--overwrite"] + ) + assert result.exit_code == 0 + files_present(output_dir) + def test_tree_plot_basic_cli_abbrev(cls): output_dir = "tree_plot_basic_cli_abbrev" result = cls.runner.invoke(