fix flake8 error

Author: aeeeeeep

objwatch/core.py

@@ -9,9 +9,9 @@
 from .config import ObjWatchConfig
 from .tracer import Tracer
 from .wrappers import ABCWrapper
-from .runtime_info import runtime_info
-from .sinks.consumer import ZeroMQFileConsumer, DynamicRoutingConsumer
+from .sinks.consumer import DynamicRoutingConsumer
 from .utils.logger import log_info, setup_logging_from_config
+from .runtime_info import runtime_info
 
 
 class ObjWatch:

objwatch/sinks/consumer.py

@@ -1,17 +1,15 @@
 # MIT License
 # Copyright (c) 2025 aeeeeeep
 
+import os
 import zmq
 import time
 import logging
 import msgpack
 import threading
-import os
-import sys
 from pathlib import Path
 from typing import Dict, Any, Optional
 from collections import OrderedDict
-from functools import lru_cache
 
 
 class ZeroMQFileConsumer:

objwatch/utils/logger.py

@@ -130,16 +130,17 @@ def emit(self, record: logging.LogRecord) -> None:
                 'time': record.created,
                 'name': record.name,
             }
-            
+
             # Add output_file and process_id for dynamic routing support
             if hasattr(sink, 'output_file') and sink.output_file:
                 event['output_file'] = sink.output_file
             else:
                 event['output_file'] = None
-            
+
             import os
+
             event['process_id'] = os.getpid()
-            
+
             sink.emit(event)
         except Exception as e:
             logging.error(f"SinkHandler failed to emit record: {e}")

tests/test_base.py

@@ -3,17 +3,16 @@
 
 import os
 import runpy
-import importlib
-import unittest
-from unittest.mock import MagicMock, patch
 import logging
+import unittest
+import importlib
 from io import StringIO
+from unittest.mock import MagicMock, patch
+
 import objwatch
-from objwatch.config import ObjWatchConfig
 from objwatch.wrappers import BaseWrapper, TensorShapeWrapper, ABCWrapper
 from objwatch.core import ObjWatch
 from objwatch.targets import Targets
-from objwatch.tracer import Tracer
 from tests.util import strip_line_numbers
 
 try:
@@ -285,7 +284,10 @@ def test_wrap_call_with_tensor_dict_over_limit(self):
 
         tensors_dict = {f"key_{i}": torch.randn(2, 2) for i in range(5)}
         mock_frame.f_locals = {'arg_tensors': tensors_dict}
-        expected_call_msg = "'0':(dict)[('key_0', torch.Size([2, 2])), ('key_1', torch.Size([2, 2])), ('key_2', torch.Size([2, 2])), '... (2 more elements)']"
+        expected_call_msg = (
+            "'0':(dict)[('key_0', torch.Size([2, 2])), ('key_1', torch.Size([2, 2])), "
+            "('key_2', torch.Size([2, 2])), '... (2 more elements)']"
+        )
         actual_call_msg = self.tensor_shape_logger.wrap_call('test_tensor_func', mock_frame)
         self.assertEqual(actual_call_msg, expected_call_msg)
 

tests/test_comprehensive_exclude.py

@@ -1,19 +1,9 @@
-#!/usr/bin/env python3
-"""
-Comprehensive test for exclude functionality in track_all mode.
-"""
-
-import sys
-import os
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+# MIT License
+# Copyright (c) 2025 aeeeeeep
 
 from objwatch.tracer import Tracer
 from objwatch.config import ObjWatchConfig
 
-# Import test module from the same directory
-from .utils.example_module import TestClass
-
 
 def test_comprehensive_exclude():
     """Test comprehensive exclude functionality with track_all mode."""

tests/test_exclude_functionality.py

@@ -1,17 +1,10 @@
-#!/usr/bin/env python3
-"""Test script to verify exclude functionality with track_all."""
+# MIT License
+# Copyright (c) 2025 aeeeeeep
 
 import sys
-import os
 
-# Add the objwatch package to the path
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from objwatch.config import ObjWatchConfig
 from objwatch.tracer import Tracer
-
-# Import test module from the same directory
-from .utils.example_module import TestClass
+from objwatch.config import ObjWatchConfig
 
 
 def test_exclude_functionality():
@@ -50,10 +43,10 @@ def test_exclude_functionality():
     print(f"Should track excluded_attr: {should_track_attr_excluded}")
 
     # Verify results
-    assert should_track_tracked == True, "tracked_method should be tracked"
-    assert should_track_excluded == False, "excluded_method should be excluded"
-    assert should_track_attr_tracked == True, "tracked_attr should be tracked"
-    assert should_track_attr_excluded == False, "excluded_attr should be excluded"
+    assert should_track_tracked, "tracked_method should be tracked"
+    assert not should_track_excluded, "excluded_method should be excluded"
+    assert should_track_attr_tracked, "tracked_attr should be tracked"
+    assert not should_track_attr_excluded, "excluded_attr should be excluded"
 
     print("All exclude functionality tests passed!")
     # All assertions passed, no return value needed for pytest

tests/test_multiprocessing_handls.py

@@ -1,3 +1,6 @@
+# MIT License
+# Copyright (c) 2025 aeeeeeep
+
 import runpy
 import unittest
 from objwatch import ObjWatch

tests/test_zmq_dynamic_routing.py

@@ -5,8 +5,9 @@
 import time
 import tempfile
 import unittest
-from objwatch.sinks.consumer import DynamicRoutingConsumer
+
 from objwatch.sinks.zmq_sink import ZeroMQSink
+from objwatch.sinks.consumer import DynamicRoutingConsumer
 
 
 class TestDynamicRoutingConsumer(unittest.TestCase):
@@ -28,16 +29,17 @@ def tearDown(self):
         """
         # Clean up test output files
         if os.path.exists(self.temp_dir):
+            import logging
             for filename in os.listdir(self.temp_dir):
                 filepath = os.path.join(self.temp_dir, filename)
                 try:
                     os.remove(filepath)
-                except Exception:
-                    pass
+                except Exception as e:
+                    logging.debug(f"Failed to remove {filepath}: {e}")
             try:
                 os.rmdir(self.temp_dir)
-            except Exception:
-                pass
+            except Exception as e:
+                logging.debug(f"Failed to remove directory {self.temp_dir}: {e}")
 
     def test_dynamic_routing_basic(self):
         """
@@ -59,7 +61,7 @@ def test_dynamic_routing_basic(self):
 
         # Give consumer time to start and connect
         # Increase delay to handle ZeroMQ SUB socket's slow joiner problem
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Send messages with different output_file
         # Send multiple messages to increase chance of reception
@@ -88,14 +90,14 @@ def test_dynamic_routing_basic(self):
                 "output_file": output1,
                 "process_id": os.getpid(),
             }
-            
+
             sink.emit(event1)
             sink.emit(event2)
             sink.emit(event3)
             time.sleep(0.1)
 
         # Give time for messages to be processed
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Clean up
         consumer.stop()
@@ -114,18 +116,18 @@ def test_dynamic_routing_basic(self):
         # Check that at least some messages were received
         self.assertTrue(len(content1) > 0, "Output file 1 should contain messages")
         self.assertTrue(len(content2) > 0, "Output file 2 should contain messages")
-        
+
         # Check for presence of expected messages (may not be all due to ZeroMQ async nature)
         if "Message to output1" in content1:
             print("✓ Received 'Message to output1'")
         else:
             print("✗ Did not receive 'Message to output1' (may be due to ZeroMQ timing)")
-            
+
         if "Another message to output1" in content1:
             print("✓ Received 'Another message to output1'")
         else:
             print("✗ Did not receive 'Another message to output1' (may be due to ZeroMQ timing)")
-            
+
         if "Message to output2" in content2:
             print("✓ Received 'Message to output2'")
         else:
@@ -172,13 +174,13 @@ def test_file_handle_lru_cache(self):
         Test LRU cache for file handles.
         """
         max_open_files = 3
-        
+
         # Create ZeroMQSink first and bind to endpoint
         sink = ZeroMQSink(endpoint=self.endpoint, topic="")
-        
+
         # Wait a bit for sink to be ready
         time.sleep(0.1)
-        
+
         consumer = DynamicRoutingConsumer(
             endpoint=self.endpoint,
             auto_start=True,
@@ -189,7 +191,7 @@ def test_file_handle_lru_cache(self):
 
         # Give consumer time to start and connect
         # Increase delay to handle ZeroMQ SUB socket's slow joiner problem
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Create more output files than max_open_files
         output_files = [os.path.join(self.temp_dir, f"output{i}.log") for i in range(5)]
@@ -209,7 +211,7 @@ def test_file_handle_lru_cache(self):
                 time.sleep(0.05)
 
         # Give time for messages to be processed
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Clean up
         consumer.stop()
@@ -226,20 +228,18 @@ def test_consumer_lifecycle(self):
         Test proper lifecycle management of DynamicRoutingConsumer.
         """
         # Create consumer
-        consumer = DynamicRoutingConsumer(
-            endpoint=self.endpoint, auto_start=False, allowed_directories=[self.temp_dir]
-        )
+        consumer = DynamicRoutingConsumer(endpoint=self.endpoint, auto_start=False, allowed_directories=[self.temp_dir])
 
         # Start consumer
         consumer.start()
-        time.sleep(0.2)
+        time.sleep(0.1)
 
         # Verify consumer is running
         self.assertTrue(consumer.running, "Consumer should be running after start()")
 
         # Stop consumer
         consumer.stop()
-        time.sleep(0.2)
+        time.sleep(0.1)
 
         # Verify consumer has stopped
         self.assertFalse(consumer.running, "Consumer should not be running after stop()")
@@ -254,7 +254,7 @@ def test_consumer_context_manager(self):
         ) as consumer:
             # Start consumer within context
             consumer.start()
-            time.sleep(0.2)
+            time.sleep(0.1)
             self.assertTrue(consumer.running, "Consumer should be running within context")
 
         # Verify consumer has been stopped after context exit
@@ -284,18 +284,18 @@ def test_process_id_in_output(self):
 
         # Create ZeroMQSink first and bind to endpoint
         sink = ZeroMQSink(endpoint=self.endpoint, topic="")
-        
+
         # Wait a bit for sink to be ready
         time.sleep(0.1)
-        
+
         # Create and start consumer
         consumer = DynamicRoutingConsumer(
             endpoint=self.endpoint, auto_start=True, daemon=True, allowed_directories=[self.temp_dir]
         )
 
         # Give consumer time to start and connect
         # Increase delay to handle ZeroMQ SUB socket's slow joiner problem
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Send multiple messages to increase chance of reception
         for _ in range(5):
@@ -311,7 +311,7 @@ def test_process_id_in_output(self):
             time.sleep(0.1)
 
         # Give time for messages to be processed
-        time.sleep(0.5)
+        time.sleep(0.1)
 
         # Clean up
         consumer.stop()
@@ -321,10 +321,10 @@ def test_process_id_in_output(self):
         if os.path.exists(output_file):
             with open(output_file, "r") as f:
                 content = f.read()
-            
+
             # Check that at least some messages were received
             self.assertTrue(len(content) > 0, "Output file should contain messages")
-            
+
             # Check for process ID (may not be present if no messages were received)
             if "PID:12345" in content:
                 print("✓ Process ID found in output")

tests/test_zmq_e2e.py

@@ -5,6 +5,7 @@
 import time
 import unittest
 import tempfile
+
 from objwatch import ObjWatch, watch
 from objwatch.config import ObjWatchConfig
 from objwatch.sinks.consumer import ZeroMQFileConsumer

tests/test_zmq_integration.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import time
+
 from objwatch import ObjWatch
 from objwatch.config import ObjWatchConfig