refactor

Author: stekern

src/tfblocks/main.py

@@ -21,10 +21,11 @@ def file_exists(file_path: str) -> bool:
     """Check if a file exists."""
     return os.path.exists(file_path)
 
+
 def extract_resource_addresses_from_content(content: str) -> List[str]:
     """Extract resource and module addresses from Terraform content."""
     addresses = []
-    
+
     # Match resource blocks
     resource_pattern = r'resource\s+"([^"]+)"\s+"([^"]+)"'
     resources = re.finditer(resource_pattern, content)
@@ -36,35 +37,39 @@ def extract_resource_addresses_from_content(content: str) -> List[str]:
     modules = re.finditer(module_pattern, content)
     for match in modules:
         addresses.append(f"module.{match.group(1)}")
-        
+
     return addresses
 
+
 def extract_resource_addresses_from_file(file_path: str) -> List[str]:
     """Extract resource and module addresses from a Terraform file."""
     addresses = []
-    
+
     # First check if file exists
     if not file_exists(file_path):
         print(f"Error: File {file_path} does not exist", file=sys.stderr)
         sys.exit(1)
-        
+
     try:
         with open(file_path, "r") as f:
             content = f.read()
-            
+
         addresses = extract_resource_addresses_from_content(content)
 
     except Exception as e:
         print(f"Warning: Could not process file {file_path}: {str(e)}", file=sys.stderr)
 
     if not addresses:
-        print(f"Warning: No resources or modules found in file {file_path}", file=sys.stderr)
-        
+        print(
+            f"Warning: No resources or modules found in file {file_path}",
+            file=sys.stderr,
+        )
+
     return addresses
 
 
 def is_resource_match(
-    resource_addr: str, filter_addrs: List[str], file_addrs: List[str]
+    resource_addr: str, filter_addrs: List[str] = [], file_addrs: List[str] = []
 ) -> bool:
     """Check if resource address matches the filter conditions.
 
@@ -106,25 +111,26 @@ def matches_address_list(addr: str, addr_list: List[str]) -> bool:
             if base_addr_match:
                 # Get the base address and any remaining part after the index
                 base_addr = base_addr_match.group(1)
-                remaining = base_addr_match.group(3)
 
                 # Check if the base address matches the filter
                 if base_addr == filter_addr:
                     return True
 
                 # Handle nested modules with indices
-                if filter_addr.startswith("module.") and base_addr.startswith(f"{filter_addr}."):
+                if filter_addr.startswith("module.") and base_addr.startswith(
+                    f"{filter_addr}."
+                ):
                     return True
 
             # Standard module prefix match (for non-indexed modules)
             if filter_addr.startswith("module.") and addr.startswith(f"{filter_addr}."):
                 return True
-                
+
             # For file addresses, we need to check if the resource type and name match
             # even if the full address is different due to modules
             resource_type, resource_name = extract_resource_type_and_name(addr)
             filter_type, filter_name = extract_resource_type_and_name(filter_addr)
-            
+
             if resource_type and resource_name and filter_type and filter_name:
                 if resource_type == filter_type and resource_name == filter_name:
                     return True
@@ -152,23 +158,25 @@ def matches_address_list(addr: str, addr_list: List[str]) -> bool:
 def filter_resources(
     state: Dict[str, Any], addresses: List[str] = [], files: List[str] = []
 ) -> List[Dict[str, Any]]:
-    """Extract matching AWS managed resources from state."""
+    """Extract matching AWS resources from Terraform state."""
     # Extract addresses from files if provided
     file_addresses = []
     if files:
         for file_path in files:
             extracted = extract_resource_addresses_from_file(file_path)
             file_addresses.extend(extracted)
-        
+
         if not file_addresses:
-            print("Warning: No resources or modules found in any of the specified files.", file=sys.stderr)
+            print(
+                "Warning: No resources or modules found in any of the specified files.",
+                file=sys.stderr,
+            )
             print("No resources will be included in the output.", file=sys.stderr)
             return []
 
     resources = []
     modules_to_process = [state["values"]["root_module"]]
 
-    # Iterative approach instead of recursion
     while modules_to_process:
         module = modules_to_process.pop()
 
@@ -185,15 +193,18 @@ def filter_resources(
                 resources.append(resource)
 
     if not resources and (addresses or files):
-        print("Warning: No resources found matching the specified filters.", file=sys.stderr)
+        print(
+            "Warning: No resources found matching the specified filters.",
+            file=sys.stderr,
+        )
 
     return resources
 
 
 def generate_import_block(
     resource: Dict[str, Any], schema_classes: Dict[str, type]
 ) -> str:
-    """Generate import block for a resource."""
+    """Generate Terraform import block for a resource."""
     matching_class = schema_classes.get(resource["type"])
     documentation = f"https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/{resource['type'].replace('aws_', '')}#import"
     import_id = f'"" # TODO: {documentation}'
@@ -213,7 +224,7 @@ def generate_import_block(
 
 
 def generate_removed_block(resource_addr: str, destroy: bool = False) -> str:
-    """Generate removed block for a resource or module."""
+    """Generate Terraform removed block for a resource or module."""
     destroy_line = "\n    destroy = true" if destroy else "\n    destroy = false"
 
     # Strip instance keys from the resource address for the 'from' attribute
@@ -227,23 +238,52 @@ def generate_removed_block(resource_addr: str, destroy: bool = False) -> str:
 }}"""
 
 
+def generate_blocks_for_command(
+    resources: List[Dict[str, Any]], command: str, destroy: bool = False
+) -> List[str]:
+    """Generate Terraform code blocks based on command."""
+    blocks = []
+
+    if command == "remove":
+        # For removed blocks, deduplicate based on the base address
+        base_addresses = set()
+
+        for resource in resources:
+            # Strip instance keys to get the base address
+            base_addr = re.sub(
+                r'(\[\d+\]|\[".+?"\]|\[\'.+?\'\])', "", resource["address"]
+            )
+
+            # Only add a removed block if we haven't seen this base address before
+            if base_addr not in base_addresses:
+                base_addresses.add(base_addr)
+                blocks.append(generate_removed_block(resource["address"], destroy))
+    elif command == "import":
+        # For import blocks, we need the full resource data
+        schema_classes = get_aws_resource_import_id_generators()
+        blocks = [generate_import_block(r, schema_classes) for r in resources]
+    else:
+        raise ValueError(f"Invalid command '{command}'")
+    return blocks
+
+
 def parse_args() -> argparse.Namespace:
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(
         description="Terraform blocks utility for generating and managing Terraform blocks",
-        epilog="Example usage: terraform show -json | tfblocks import [resource_addresses]"
+        epilog="Example usage: terraform show -json | tfblocks import [resource_addresses]",
     )
 
-    # Global options
     parser.add_argument(
         "--no-color", action="store_true", help="Disable colored output"
     )
 
-    # Create subcommands
-    subparsers = parser.add_subparsers(dest="command", help="Command to execute", required=True)
+    subparsers = parser.add_subparsers(
+        dest="command", help="Command to execute", required=True
+    )
 
     # Common filter arguments function
-    def add_filter_args(cmd_parser):
+    def add_filter_args(cmd_parser: argparse.ArgumentParser):
         cmd_parser.add_argument(
             "addresses",
             nargs="*",
@@ -253,11 +293,9 @@ def add_filter_args(cmd_parser):
             "--files", "-f", nargs="+", help="Optional Terraform files to filter by"
         )
 
-    # Import command
     import_parser = subparsers.add_parser("import", help="Generate import blocks")
     add_filter_args(import_parser)
 
-    # Remove command
     remove_parser = subparsers.add_parser("remove", help="Generate removed blocks")
     add_filter_args(remove_parser)
     remove_parser.add_argument(
@@ -266,73 +304,44 @@ def add_filter_args(cmd_parser):
         help="Set destroy = true in removed blocks (default is false)",
     )
 
-    # List command
-    list_parser = subparsers.add_parser("list", help="List resource addresses")
+    list_parser = subparsers.add_parser(
+        "list", help="List addresses delimited by newlines"
+    )
     add_filter_args(list_parser)
 
-    # No default command - require explicit command
     args = parser.parse_args()
     return args
 
 
-def generate_blocks_for_command(
-    resources: List[Dict[str, Any]], command: str, destroy: bool = False
-) -> List[str]:
-    """Generate appropriate blocks based on command."""
-    blocks = []
-
-    if command == "remove":
-        # For removed blocks, deduplicate based on the base address
-        base_addresses = set()
-
-        for resource in resources:
-            # Strip instance keys to get the base address
-            base_addr = re.sub(
-                r'(\[\d+\]|\[".+?"\]|\[\'.+?\'\])', "", resource["address"]
-            )
-
-            # Only add a removed block if we haven't seen this base address before
-            if base_addr not in base_addresses:
-                base_addresses.add(base_addr)
-                blocks.append(generate_removed_block(resource["address"], destroy))
-    else:  # import command
-        # For import blocks, we need the full resource data
-        schema_classes = get_aws_resource_import_id_generators()
-        blocks = [generate_import_block(r, schema_classes) for r in resources]
-
-    return blocks
-
-
 def main():
     args = parse_args()
-    
     state = json.load(sys.stdin)
-
     if state.get("format_version") != "1.0":
         print(
             "Error: Unsupported state file format version. Expected version '1.0'.",
             file=sys.stderr,
         )
         sys.exit(1)
-
     resources = filter_resources(state, args.addresses, args.files)
 
     if args.command == "list":
-        # Just output the addresses
         for resource in resources:
             print(resource["address"])
         return
 
-    # Generate blocks based on command
     blocks = generate_blocks_for_command(
         resources, args.command, getattr(args, "destroy", False)
     )
 
-    # Output results
     if args.no_color:
         print("\n\n".join(blocks))
     else:
-        print("\033[92m" + "\n\n".join(blocks) + "\033[0m")
+        if args.command == "import":
+            print("\033[92m" + "\n\n".join(blocks) + "\033[0m")
+        elif args.command == "remove":
+            print("\033[91m" + "\n\n".join(blocks) + "\033[0m")
+        else:
+            print("\033[92m" + "\n\n".join(blocks) + "\033[0m")
 
 
 if __name__ == "__main__":