fix linting

Author: KeplerC

robodm/agent/executor.py

@@ -128,7 +128,7 @@ def ray_map_wrapper(batch):
                     batch_dict = batch
 
                 batch_size = len(next(iter(batch_dict.values())))
-                transformed_batch = {}
+                transformed_batch: Dict[str, List[Any]] = {}
 
                 for i in range(batch_size):
                     # Extract single trajectory from batch

robodm/agent/tools/config.py

@@ -5,7 +5,7 @@
 and helper functions for creating custom configurations.
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 
 def create_vision_config(model: str = "qwen2.5-7b",
@@ -108,7 +108,7 @@ def create_custom_config(
     Returns:
         Custom configuration dictionary
     """
-    config = {}
+    config: Dict[str, Any] = {}
 
     if tool_parameters:
         config["tools"] = tool_parameters
@@ -232,7 +232,7 @@ def merge_configs(*configs: Dict[str, Any]) -> Dict[str, Any]:
     Returns:
         Merged configuration dictionary
     """
-    result = {}
+    result: Dict[str, Any] = {}
 
     for config in configs:
         if not isinstance(config, dict):
@@ -289,7 +289,7 @@ def get_default_config() -> Dict[str, Any]:
 
 
 # Configuration presets for common scenarios
-PRESET_CONFIGS = {
+PRESET_CONFIGS: Dict[str, Callable[..., Dict[str, Any]]] = {
     "vision": create_vision_config,
     "analysis": create_analysis_config,
     "minimal": create_minimal_config,

robodm/backend/base.py

@@ -13,7 +13,7 @@ class StreamMetadata:
     feature_type: str  # Using string to avoid circular imports with FeatureType
     encoding: str
     time_base: tuple[int, int]  # Numerator, denominator for time base fraction
-    additional_metadata: Dict[str, str] = None
+    additional_metadata: Optional[Dict[str, str]] = None
 
 
 @dataclass
@@ -176,7 +176,7 @@ def seek_container(self,
     @abstractmethod
     def decode_stream_frames(self,
                              stream_index: int,
-                             packet_data: bytes = None) -> List[Any]:
+                             packet_data: Optional[bytes] = None) -> List[Any]:
         """Decode frames from a stream, optionally with packet data
 
         Args:

robodm/backend/codec_config.py

@@ -33,9 +33,9 @@ def is_codec_config_supported(width: int,
         """Check if a specific width/height/pixel format combination is supported by codec."""
         try:
             cc = av.codec.CodecContext.create(codec_name, "w")
-            cc.width = width
-            cc.height = height
-            cc.pix_fmt = pix_fmt
+            cc.width = width  # type: ignore[attr-defined]
+            cc.height = height  # type: ignore[attr-defined]
+            cc.pix_fmt = pix_fmt  # type: ignore[attr-defined]
             cc.time_base = Fraction(1, 30)
             cc.open(strict=True)
             # Note: CodecContext doesn't have a close() method in newer PyAV versions
@@ -92,7 +92,7 @@ def is_raw_data_codec(codec_name: str) -> bool:
         return codec_name.startswith("rawvideo")
 
     # Image codec configurations (use actual codec for container)
-    IMAGE_CODEC_CONFIGS = {
+    IMAGE_CODEC_CONFIGS: Dict[str, Dict[str, Any]] = {
         "libx264": {
             "container_codec": "libx264",  # Use actual codec for container
             "pixel_format": "yuv420p",
@@ -126,7 +126,7 @@ def is_raw_data_codec(codec_name: str) -> bool:
     }
 
     # Raw data codec configurations (always use rawvideo container)
-    RAW_DATA_CODEC_CONFIGS = {
+    RAW_DATA_CODEC_CONFIGS: Dict[str, Dict[str, Any]] = {
         "rawvideo": {
             "container_codec": "rawvideo",  # Always rawvideo for container
             "internal_codec": "pickle_raw",  # Default internal implementation
@@ -279,15 +279,16 @@ def get_codec_for_feature(self,
         is_rgb_image = (data_shape is not None and len(data_shape) == 3
                         and data_shape[2] == 3)
 
-        if is_rgb_image:
+        if is_rgb_image and data_shape is not None:
             # This is RGB image data - can use video codecs
             height, width = data_shape[0], data_shape[1]
 
             # Check if a specific video codec was provided
             if self.video_codec and self.video_codec != "auto":
-                if self.is_image_codec(
-                        self.video_codec) and self.is_valid_image_shape(
-                            data_shape, self.video_codec):
+                if (self.is_image_codec(self.video_codec)
+                        and data_shape is not None
+                        and self.is_valid_image_shape(data_shape,
+                                                      self.video_codec)):
                     logger.debug(
                         f"Using specified video codec {self.video_codec} for RGB shape {data_shape}"
                     )
@@ -299,7 +300,8 @@ def get_codec_for_feature(self,
 
             # Check if user specified a general codec other than auto
             if self.codec != "auto" and self.is_image_codec(self.codec):
-                if self.is_valid_image_shape(data_shape, self.codec):
+                if data_shape is not None and self.is_valid_image_shape(
+                        data_shape, self.codec):
                     logger.debug(
                         f"Using user-specified image codec {self.codec} for RGB shape {data_shape}"
                     )
@@ -318,7 +320,8 @@ def get_codec_for_feature(self,
             ]
 
             for codec in codec_preferences:
-                if self.is_valid_image_shape(data_shape, codec):
+                if data_shape is not None and self.is_valid_image_shape(
+                        data_shape, codec):
                     logger.debug(
                         f"Selected image codec {codec} for RGB shape {data_shape}"
                     )
@@ -377,16 +380,19 @@ def _can_codec_handle_feature(self, codec: str,
     def get_container_codec(self, codec: str) -> str:
         """Get the container codec name for a given codec."""
         if codec in self.IMAGE_CODEC_CONFIGS:
-            return self.IMAGE_CODEC_CONFIGS[codec]["container_codec"]
+            return cast(str,
+                        self.IMAGE_CODEC_CONFIGS[codec]["container_codec"])
         elif codec in self.RAW_DATA_CODEC_CONFIGS:
-            return self.RAW_DATA_CODEC_CONFIGS[codec]["container_codec"]
+            return cast(str,
+                        self.RAW_DATA_CODEC_CONFIGS[codec]["container_codec"])
         else:
             raise ValueError(f"Unknown codec {codec}")
 
     def get_internal_codec(self, codec: str) -> Optional[str]:
         """Get the internal codec implementation name for raw data codecs."""
         if codec in self.RAW_DATA_CODEC_CONFIGS:
-            return self.RAW_DATA_CODEC_CONFIGS[codec]["internal_codec"]
+            return cast(str,
+                        self.RAW_DATA_CODEC_CONFIGS[codec]["internal_codec"])
         elif codec in self.IMAGE_CODEC_CONFIGS:
             # Image codecs don't have internal codecs
             return None
@@ -402,15 +408,18 @@ def get_raw_codec_name(self, codec: str) -> str:
         # Fallback for backward compatibility
         legacy_configs = self.CODEC_CONFIGS
         if codec in legacy_configs:
-            return legacy_configs[codec].get("raw_codec", "pickle_raw")
+            return cast(str, legacy_configs[codec].get("raw_codec",
+                                                       "pickle_raw"))
 
         return "pickle_raw"
 
     def get_pixel_format(self, codec: str,
                          feature_type: FeatureType) -> Optional[str]:
         """Get appropriate pixel format for codec and feature type."""
         if codec in self.IMAGE_CODEC_CONFIGS:
-            base_format = self.IMAGE_CODEC_CONFIGS[codec].get("pixel_format")
+            base_format = cast(
+                Optional[str],
+                self.IMAGE_CODEC_CONFIGS[codec].get("pixel_format"))
 
             # For FFV1, use RGB24 to avoid YUV conversion issues
             if codec == "ffv1":
@@ -434,14 +443,15 @@ def get_codec_options(self, codec: str) -> Dict[str, Any]:
 
         if codec in self.IMAGE_CODEC_CONFIGS:
             # Video/image codec - only use video-specific options
-            default_options = self.IMAGE_CODEC_CONFIGS[codec].get(
-                "options", {}).copy()
+            options_dict = self.IMAGE_CODEC_CONFIGS[codec].get("options", {})
+            default_options = cast(Dict[str, Any], options_dict).copy()
             # Only merge video-specific custom options
             default_options.update(self.video_custom_options)
         elif codec in self.RAW_DATA_CODEC_CONFIGS:
             # Raw data codec - only use raw-specific options
-            default_options = (self.RAW_DATA_CODEC_CONFIGS[codec].get(
-                "options", {}).copy())
+            options_dict = self.RAW_DATA_CODEC_CONFIGS[codec].get(
+                "options", {})
+            default_options = cast(Dict[str, Any], options_dict).copy()
             # Only merge raw-specific custom options
             default_options.update(self.raw_custom_options)
 
@@ -451,7 +461,7 @@ def get_codec_options(self, codec: str) -> Dict[str, Any]:
     def for_transcoding_to_internal_codec(
             cls,
             internal_codec: str,
-            codec_options: Optional[Dict[str, Any]] = None) -> "CodecConfig":
+            codec_options: Optional[Dict[str, Any]] = None) -> Any:
         """Create a CodecConfig specifically for transcoding to a particular internal codec.
 
         This is used during transcoding operations where we need to convert between

robodm/backend/codecs.py

@@ -231,7 +231,7 @@ def get_container_encoding(self) -> str:
 class PyAVVideoCodec(VideoCodec):
     """PyAV-based video codec wrapper"""
 
-    def __init__(self, codec_name: str = None, **kwargs):
+    def __init__(self, codec_name: Optional[str] = None, **kwargs):
         # Handle both old and new initialization styles
         if codec_name is None:
             # New style: codec name should be passed as kwarg or inferred from registration