diff --git a/struct_module/commands/generate.py b/struct_module/commands/generate.py index fad546b..32103a6 100644 --- a/struct_module/commands/generate.py +++ b/struct_module/commands/generate.py @@ -39,19 +39,21 @@ def execute(self, args): def _create_structure(self, args): if isinstance(args, dict): args = argparse.Namespace(**args) + this_file = os.path.dirname(os.path.realpath(__file__)) + contribs_path = os.path.join(this_file, "..", "contribs") + if args.structure_definition.startswith("file://") and args.structure_definition.endswith(".yaml"): with open(args.structure_definition[7:], 'r') as f: config = yaml.safe_load(f) else: - if args.structures_path is None: - this_file = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(this_file, "..", "contribs", f"{args.structure_definition}.yaml") - else: + file_path = os.path.join(contribs_path, f"{args.structure_definition}.yaml") + if args.structures_path: file_path = os.path.join(args.structures_path, f"{args.structure_definition}.yaml") - # show error if file is not found + if not os.path.exists(file_path): self.logger.error(f"File not found: {file_path}") return + with open(file_path, 'r') as f: config = yaml.safe_load(f) diff --git a/struct_module/commands/list.py b/struct_module/commands/list.py index 1609139..0e20546 100644 --- a/struct_module/commands/list.py +++ b/struct_module/commands/list.py @@ -16,15 +16,22 @@ def execute(self, args): self._list_structures(args) def _list_structures(self, args): + this_file = os.path.dirname(os.path.realpath(__file__)) + contribs_path = os.path.join(this_file, "..", "contribs") - if args.structures_path is None: - this_file = os.path.dirname(os.path.realpath(__file__)) - final_path = os.path.join(this_file, "..", "contribs") + if args.structures_path: + final_path = args.structures_path + paths_to_list = [final_path, contribs_path] else: - final_path = os.path.join(args.structures_path) + paths_to_list = [contribs_path] print("Listing available structures") - sorted_list = [structure for structure in os.listdir(final_path) if structure.endswith('.yaml')] - sorted_list.sort() + all_structures = set() + for path in paths_to_list: + if os.path.exists(path): + structures = [structure for structure in os.listdir(path) if structure.endswith('.yaml')] + all_structures.update(structures) + + sorted_list = sorted(all_structures) for structure in sorted_list: print(f" - {structure[:-5]}")