Fix formatting

Signed-off-by: Victor Chang <[email protected]>

Author: mocsharp

tools/pipeline-generator/pipeline_generator/cli/main.py

@@ -99,7 +99,7 @@ def list(ctx: click.Context, format: str, bundles_only: bool, tested_only: bool,
 
     # Fetch models from all endpoints
     console.print("[blue]Fetching models from HuggingFace...[/blue]")
-    
+
     if all:
         # Show all models, but fetch detailed info for MONAI Bundles to get accurate extension data
         console.print("[yellow]Note: Fetching detailed info for MONAI Bundles to show accurate extension data[/yellow]")
@@ -263,7 +263,7 @@ def _display_table(models: List[ModelInfo], verified_models: Set[str]) -> None:
                 bundle_status = f"[dim]✗ No ({primary_ext})[/dim]"
             else:
                 bundle_status = "[dim]✗ No[/dim]"
-        
+
         status = "[bold green]✓ Verified[/bold green]" if model.model_id in verified_models else ""
         table.add_row(
             model.model_id,

tools/pipeline-generator/pipeline_generator/config/settings.py

@@ -55,8 +55,7 @@ class Settings(BaseModel):
     """Application settings."""
 
     supported_models: List[str] = Field(
-        default_factory=lambda: [".ts", ".pt", ".pth", ".safetensors"],
-        description="Supported model file extensions"
+        default_factory=lambda: [".ts", ".pt", ".pth", ".safetensors"], description="Supported model file extensions"
     )
     endpoints: List[Endpoint] = Field(default_factory=list)
     additional_models: List[Endpoint] = Field(default_factory=list)

tools/pipeline-generator/pipeline_generator/core/hub_client.py

@@ -28,15 +28,13 @@ class HuggingFaceClient:
 
     def __init__(self, settings: Optional[Settings] = None) -> None:
         """Initialize the HuggingFace Hub client.
-        
+
         Args:
             settings: Pipeline generator settings containing supported extensions
         """
         self.api = HfApi()
         self.settings = settings
-        self.supported_extensions = (
-            settings.supported_models if settings else [".ts", ".pt", ".pth", ".safetensors"]
-        )
+        self.supported_extensions = settings.supported_models if settings else [".ts", ".pt", ".pth", ".safetensors"]
 
     def list_models_from_organization(self, organization: str) -> List[ModelInfo]:
         """List all models from a HuggingFace organization.
@@ -79,7 +77,9 @@ def get_model_info(self, model_id: str) -> Optional[ModelInfo]:
             logger.error(f"Error fetching model {model_id}: {e}")
             return None
 
-    def list_models_from_endpoints(self, endpoints: List[Endpoint], fetch_details_for_bundles: bool = False) -> List[ModelInfo]:
+    def list_models_from_endpoints(
+        self, endpoints: List[Endpoint], fetch_details_for_bundles: bool = False
+    ) -> List[ModelInfo]:
         """List models from all configured endpoints.
 
         Args:
@@ -96,7 +96,7 @@ def list_models_from_endpoints(self, endpoints: List[Endpoint], fetch_details_fo
                 # List all models from organization
                 logger.info(f"Fetching models from organization: {endpoint.organization}")
                 models = self.list_models_from_organization(endpoint.organization)
-                
+
                 if fetch_details_for_bundles:
                     # For MONAI organization, fetch detailed info for all models to get accurate extension data
                     # This is needed because bulk API doesn't provide file information (siblings=None)
@@ -118,7 +118,7 @@ def list_models_from_endpoints(self, endpoints: List[Endpoint], fetch_details_fo
                             else:
                                 enhanced_models.append(model)
                         models = enhanced_models
-                
+
                 all_models.extend(models)
 
             elif endpoint.model_id:
@@ -140,7 +140,7 @@ def _detect_model_extensions(self, model_data: Any) -> List[str]:
             List of detected model file extensions
         """
         extensions = []
-        
+
         try:
             if hasattr(model_data, "siblings") and model_data.siblings is not None:
                 file_names = [f.rfilename for f in model_data.siblings]
@@ -151,13 +151,13 @@ def _detect_model_extensions(self, model_data: Any) -> List[str]:
                                 extensions.append(ext)
         except Exception as e:
             logger.debug(f"Could not detect extensions for {getattr(model_data, 'modelId', 'unknown')}: {e}")
-        
+
         return extensions
 
     def list_torchscript_models(self, endpoints: List[Endpoint]) -> List[ModelInfo]:
         """List models that have TorchScript (.ts) files.
 
-        This method fetches detailed information for each model individually to 
+        This method fetches detailed information for each model individually to
         check for TorchScript files, which is slower than bulk listing but accurate.
 
         Args:
@@ -167,7 +167,7 @@ def list_torchscript_models(self, endpoints: List[Endpoint]) -> List[ModelInfo]:
             List of ModelInfo objects that contain .ts files
         """
         torchscript_models = []
-        
+
         for endpoint in endpoints:
             if endpoint.organization:
                 # List all models from organization first (bulk)

tools/pipeline-generator/pipeline_generator/core/models.py

@@ -80,9 +80,9 @@ def primary_extension(self) -> Optional[str]:
         """
         if not self.model_extensions:
             return None
-        
+
         # Prioritize .ts extension
         if ".ts" in self.model_extensions:
             return ".ts"
-        
+
         return self.model_extensions[0]

tools/pipeline-generator/tests/test_cli.py

@@ -250,7 +250,9 @@ def test_list_command_tested_only(self, mock_load_config, mock_client_class):
 
         # Mock the list response
         test_models = [
-            ModelInfo(model_id="MONAI/tested_model", name="Tested Model", is_monai_bundle=True, model_extensions=[".ts"]),
+            ModelInfo(
+                model_id="MONAI/tested_model", name="Tested Model", is_monai_bundle=True, model_extensions=[".ts"]
+            ),
             ModelInfo(
                 model_id="MONAI/untested_model",
                 name="Untested Model",
@@ -309,7 +311,7 @@ def test_list_command_all_flag(self, mock_load_config, mock_client_class):
                 is_monai_bundle=True,
             ),
             ModelInfo(
-                model_id="MONAI/pt_model", 
+                model_id="MONAI/pt_model",
                 name="PyTorch Model",
                 model_extensions=[".pt"],
                 is_monai_bundle=False,
@@ -331,7 +333,7 @@ def test_list_command_all_flag(self, mock_load_config, mock_client_class):
         assert "Verified: 0" in result.output
 
     @patch("pipeline_generator.cli.main.HuggingFaceClient")
-    @patch("pipeline_generator.cli.main.load_config") 
+    @patch("pipeline_generator.cli.main.load_config")
     def test_list_command_default_torchscript_only(self, mock_load_config, mock_client_class):
         """Test list command defaults to torchscript models only."""
         # Mock setup
@@ -347,7 +349,7 @@ def test_list_command_default_torchscript_only(self, mock_load_config, mock_clie
         test_models = [
             ModelInfo(
                 model_id="MONAI/ts_model",
-                name="TorchScript Model", 
+                name="TorchScript Model",
                 model_extensions=[".ts"],
                 is_monai_bundle=True,
             ),
@@ -358,10 +360,8 @@ def test_list_command_default_torchscript_only(self, mock_load_config, mock_clie
         result = self.runner.invoke(cli, ["list"])
 
         assert result.exit_code == 0
-        # Should call list_torchscript_models (default behavior) 
-        mock_client.list_torchscript_models.assert_called_once_with(
-            mock_settings.get_all_endpoints.return_value
-        )
+        # Should call list_torchscript_models (default behavior)
+        mock_client.list_torchscript_models.assert_called_once_with(mock_settings.get_all_endpoints.return_value)
         mock_client.list_models_from_endpoints.assert_not_called()
         assert "MONAI/ts_model" in result.output
         assert "Verified: 0" in result.output
@@ -389,7 +389,7 @@ def test_list_command_monai_bundle_column_logic(self, mock_load_config, mock_cli
             ),
             ModelInfo(
                 model_id="MONAI/pt_model",
-                name="PyTorch Model", 
+                name="PyTorch Model",
                 model_extensions=[".pt"],
                 is_monai_bundle=False,  # Should be False for .pt files
             ),
@@ -408,17 +408,17 @@ def test_list_command_monai_bundle_column_logic(self, mock_load_config, mock_cli
         assert result.exit_code == 0
         # Check MONAI Bundle column contents with new display format
         output = result.output
-        
+
         # Should show "✓" and "Yes" for .ts model (MONAI Bundle) - may be on separate lines due to table wrapping
         assert "MONAI/ts_model" in output
         assert "✓" in output  # Checkmark emoji
         assert "Yes" in output  # Text
-        
-        # Should show "✗" and "No" for .pt model  
+
+        # Should show "✗" and "No" for .pt model
         assert "MONAI/pt_model" in output
         assert "✗" in output  # X emoji
         assert "No" in output  # The "No" text should appear
-        
+
         # Should show "✗ No" for model with no extensions
         assert "MONAI/no_ext_model" in output
         # The Verified count may have color codes, so check for the text parts
@@ -432,10 +432,12 @@ def test_list_command_with_verified_models(self, mock_load_config, mock_client_c
         mock_settings = Mock()
         mock_settings.get_all_endpoints.return_value = [Mock(organization="MONAI")]
         mock_settings.endpoints = [
-            Mock(models=[
-                Mock(model_id="MONAI/verified_model1"),
-                Mock(model_id="MONAI/verified_model2"),
-            ])
+            Mock(
+                models=[
+                    Mock(model_id="MONAI/verified_model1"),
+                    Mock(model_id="MONAI/verified_model2"),
+                ]
+            )
         ]
         mock_load_config.return_value = mock_settings
 
@@ -466,13 +468,13 @@ def test_list_command_with_verified_models(self, mock_load_config, mock_client_c
         assert "MONAI/verified_model1" in result.output
         assert "MONAI/unverified_model" in result.output
         assert "Verified: 1" in result.output
-        
+
         # Check that verified model shows verification checkmark
-        output_lines = result.output.split('\n')
+        output_lines = result.output.split("\n")
         verified_line = [line for line in output_lines if "MONAI/verified_model1" in line]
         assert any("✓ Verified" in line for line in verified_line)
-        
-        unverified_line = [line for line in output_lines if "MONAI/unverified_model" in line] 
+
+        unverified_line = [line for line in output_lines if "MONAI/unverified_model" in line]
         assert not any("✓ Verified" in line for line in unverified_line)
 
     @patch("pipeline_generator.cli.main.HuggingFaceClient")
@@ -482,9 +484,7 @@ def test_list_command_json_output(self, mock_load_config, mock_client_class):
         # Mock setup
         mock_settings = Mock()
         mock_settings.get_all_endpoints.return_value = [Mock(organization="MONAI")]
-        mock_settings.endpoints = [
-            Mock(models=[Mock(model_id="MONAI/test_model")])
-        ]
+        mock_settings.endpoints = [Mock(models=[Mock(model_id="MONAI/test_model")])]
         mock_load_config.return_value = mock_settings
 
         mock_client = Mock()
@@ -508,17 +508,18 @@ def test_list_command_json_output(self, mock_load_config, mock_client_class):
         result = self.runner.invoke(cli, ["list", "--format", "json"])
 
         assert result.exit_code == 0
-        
+
         # Parse JSON output to verify new fields
         import json
-        json_start = result.output.find('[')
-        json_end = result.output.rfind(']') + 1  # Find the last ] and include it
+
+        json_start = result.output.find("[")
+        json_end = result.output.rfind("]") + 1  # Find the last ] and include it
         json_text = result.output[json_start:json_end]
         json_data = json.loads(json_text)
-        
+
         assert len(json_data) == 1
         model_data = json_data[0]
-        
+
         # Check all new fields are present
         assert model_data["model_id"] == "MONAI/test_model"
         assert model_data["is_monai_bundle"] is True

tools/pipeline-generator/tests/test_hub_client.py

@@ -45,10 +45,7 @@ def test_list_models_from_organization_success(self, mock_list_models):
             created_at=datetime(2023, 1, 1),
             lastModified=datetime(2023, 12, 1),
             tags=["medical", "segmentation"],
-            siblings=[
-                Mock(rfilename="configs/metadata.json"),
-                Mock(rfilename="models/model.ts")
-            ],
+            siblings=[Mock(rfilename="configs/metadata.json"), Mock(rfilename="models/model.ts")],
         )
 
         mock_model2 = SimpleModelData(
@@ -104,10 +101,7 @@ def test_get_model_info_success(self, mock_model_info):
             created_at=datetime(2023, 1, 1),
             lastModified=datetime(2023, 12, 1),
             tags=["medical", "segmentation"],
-            siblings=[
-                Mock(rfilename="configs/metadata.json"),
-                Mock(rfilename="models/model.ts")
-            ],
+            siblings=[Mock(rfilename="configs/metadata.json"), Mock(rfilename="models/model.ts")],
             cardData={"description": "Spleen segmentation model"},
         )
 
@@ -201,10 +195,7 @@ def test_extract_model_info_bundle_detection(self):
         assert model.is_monai_bundle is True
 
         # Test without TorchScript (.ts) file - only .pt file
-        mock_model.siblings = [
-            Mock(rfilename="configs/metadata.json"),
-            Mock(rfilename="models/model.pt")
-        ]
+        mock_model.siblings = [Mock(rfilename="configs/metadata.json"), Mock(rfilename="models/model.pt")]
         model = self.client._extract_model_info(mock_model)
         assert model.is_monai_bundle is False
 
@@ -384,11 +375,11 @@ def test_detect_model_extensions_with_torchscript(self):
                 Mock(rfilename="model.ts"),
                 Mock(rfilename="config.yaml"),
                 Mock(rfilename="README.md"),
-            ]
+            ],
         )
 
         extensions = self.client._detect_model_extensions(mock_model)
-        
+
         assert ".ts" in extensions
         assert len(extensions) == 1
 
@@ -402,11 +393,11 @@ def test_detect_model_extensions_multiple_formats(self):
                 Mock(rfilename="model.pt"),
                 Mock(rfilename="model.safetensors"),
                 Mock(rfilename="config.yaml"),
-            ]
+            ],
         )
 
         extensions = self.client._detect_model_extensions(mock_model)
-        
+
         assert ".ts" in extensions
         assert ".pt" in extensions
         assert ".safetensors" in extensions
@@ -420,11 +411,11 @@ def test_detect_model_extensions_no_model_files(self):
             siblings=[
                 Mock(rfilename="README.md"),
                 Mock(rfilename="config.yaml"),
-            ]
+            ],
         )
 
         extensions = self.client._detect_model_extensions(mock_model)
-        
+
         assert len(extensions) == 0
 
     def test_detect_model_extensions_no_siblings(self):
@@ -433,7 +424,7 @@ def test_detect_model_extensions_no_siblings(self):
         mock_model = SimpleModelData(modelId="MONAI/test_model")
 
         extensions = self.client._detect_model_extensions(mock_model)
-        
+
         assert len(extensions) == 0
 
     @patch("pipeline_generator.core.hub_client.list_models")
@@ -452,7 +443,7 @@ def test_list_torchscript_models_filters_correctly(self, mock_list_models):
         )
 
         mock_model_without_ts = SimpleModelData(
-            modelId="MONAI/model_without_ts", 
+            modelId="MONAI/model_without_ts",
             author="MONAI",
             downloads=50,
             likes=5,
@@ -472,13 +463,14 @@ def mock_get_model_info(model_id):
                 return self.client._extract_model_info(mock_model_without_ts)
             return None
 
-        with patch.object(self.client, 'get_model_info', side_effect=mock_get_model_info):
+        with patch.object(self.client, "get_model_info", side_effect=mock_get_model_info):
             from pipeline_generator.config.settings import Endpoint
+
             endpoints = [Endpoint(organization="MONAI")]
 
             # Test the torchscript filtering
             torchscript_models = self.client.list_torchscript_models(endpoints)
-            
+
             # Should only return the model with .ts file
             assert len(torchscript_models) == 1
             assert torchscript_models[0].model_id == "MONAI/model_with_ts"
@@ -502,7 +494,7 @@ def test_extract_model_info_includes_extensions(self):
         )
 
         result = self.client._extract_model_info(mock_model)
-        
+
         assert ".ts" in result.model_extensions
         assert ".pt" in result.model_extensions
         assert result.has_torchscript is True