fixes

Author: idanasulin2706

.gitignore

@@ -2,4 +2,6 @@ venv/
 __pycache__/
 *.py[cod]
 *$py.class
-superclient.egg-info
+superclient.egg-info
+build/
+dist/

MANIFEST.in

@@ -5,7 +5,7 @@
 
 </div>
 
-# Superstream Client For Python
+# Superclient Python
 
 A Python library for automatically optimizing Kafka producer configurations based on topic-specific recommendations.
 
@@ -82,26 +82,35 @@ The Superstream library needs to modify your producer's configuration to apply o
 ### Step 1: Install Superclient
 
 ```bash
+# Step 1: Install the package
 pip install superclient
-```
 
-### Step 2: Run
+# Step 2: One-time setup (enables automatic loading)
+python -m superclient install_pth
+```
 
-The package ships with a `sitecustomize.py` entry-point, therefore Python imports the agent automatically before your application's code starts. This is the recommended and default way to use Superclient.
+That's it! Superclient will now automatically load and optimize all Kafka producers in your Python environment.
 
-#### Manual Initialization (Only if needed)
+## Usage
 
-If `sitecustomize` is disabled in your environment (e.g., when using `python -S` or when `PYTHONNOUSERSITE` is set), you can initialize manually by adding this import at the very beginning of your application's main entry point (e.g., `main.py`, `app.py`, or `__init__.py`):
+After installation, superclient works automatically. Just use your Kafka clients as usual:
 
 ```python
-import superclient  # side-effects automatically enable the agent
-
-# Your application code follows
+# kafka-python
 from kafka import KafkaProducer
-# ... rest of your imports and code
-```
+producer = KafkaProducer(bootstrap_servers='localhost:9092')
+# Automatically optimized!
 
-Note: The manual import must be placed before any Kafka-related imports to ensure proper interception of producer creation.
+# confluent-kafka
+from confluent_kafka import Producer
+producer = Producer({'bootstrap.servers': 'localhost:9092'})
+# Automatically optimized!
+
+# aiokafka
+from aiokafka import AIOKafkaProducer
+producer = AIOKafkaProducer(bootstrap_servers='localhost:9092')
+# Automatically optimized!
+```
 
 ### Docker Integration
 
@@ -113,6 +122,9 @@ FROM python:3.8-slim
 # Install superclient
 RUN pip install superclient
 
+# Run the one-time setup to enable automatic loading
+RUN python -m superclient install_pth
+
 # Your application code
 COPY . /app
 WORKDIR /app
@@ -137,17 +149,6 @@ export SUPERSTREAM_TOPICS_LIST=orders,payments,user-events
 export SUPERSTREAM_LATENCY_SENSITIVE=true
 ```
 
-### SUPERSTREAM_LATENCY_SENSITIVE Explained
-
-The linger.ms parameter follows these rules:
-
-1. If SUPERSTREAM_LATENCY_SENSITIVE is set to true:
-   - Linger value will never be modified, regardless of other settings
-
-2. If SUPERSTREAM_LATENCY_SENSITIVE is set to false or not set:
-   - If no explicit linger exists in original configuration: Use Superstream's optimized value
-   - If explicit linger exists: Use the maximum of original value and Superstream's optimized value
-
 ## Prerequisites
 
 - Python 3.8 or higher
@@ -156,4 +157,4 @@ The linger.ms parameter follows these rules:
 
 ## License
 
-This project is licensed under the Apache License 2.0. 
+Apache License 2.0 

README.md

@@ -21,5 +21,9 @@ classifiers = [
 [project.urls]
 Homepage = "https://github.com/superstreamlabs/superclient-python"
 
+[tool.setuptools]
+packages = ["superclient", "superclient.agent", "superclient.core", "superclient.model", "superclient.util"]
+include-package-data = true
+
 [tool.setuptools.package-data]
-"superclient" = ["py.typed"] 
+"superclient" = ["py.typed", "*.pth"] 

pyproject.toml

@@ -0,0 +1,8 @@
+"""Allow running superclient module utilities"""
+import sys
+
+if len(sys.argv) > 1 and sys.argv[1] == "install_pth":
+    from .install_pth import install_pth
+    install_pth()
+else:
+    print("Usage: python -m superclient install_pth") 

sitecustomize.py

@@ -34,11 +34,20 @@ def _patch_module(module_name: str) -> None:
     """Patch a specific module if it exists in sys.modules."""
     try:
         if module_name == "kafka" and "kafka" in sys.modules:
-            patch_kafka_python(sys.modules["kafka"])
+            # Check if KafkaProducer exists before patching
+            kafka_module = sys.modules["kafka"]
+            if hasattr(kafka_module, "KafkaProducer"):
+                patch_kafka_python(kafka_module)
         elif module_name == "aiokafka" and "aiokafka" in sys.modules:
-            patch_aiokafka(sys.modules["aiokafka"])
+            # Check if AIOKafkaProducer exists before patching
+            aiokafka_module = sys.modules["aiokafka"]
+            if hasattr(aiokafka_module, "AIOKafkaProducer"):
+                patch_aiokafka(aiokafka_module)
         elif module_name == "confluent_kafka" and "confluent_kafka" in sys.modules:
-            patch_confluent(sys.modules["confluent_kafka"])
+            # Check if Producer exists before patching
+            confluent_module = sys.modules["confluent_kafka"]
+            if hasattr(confluent_module, "Producer"):
+                patch_confluent(confluent_module)
     except Exception as exc:
         logger.error("[ERR-001] Failed to patch {}: {}", module_name, exc)
 
@@ -81,7 +90,7 @@ def initialize():
     
     This function:
     1. Installs the import hook to catch future imports
-    2. Patches any pre-imported modules
+    2. Schedules patching of any pre-imported modules
     3. Starts the heartbeat thread
     """
     if is_disabled():
@@ -91,9 +100,17 @@ def initialize():
     if builtins.__import__ is not _import_hook:
         builtins.__import__ = _import_hook
 
-    # Patch any pre-imported modules
-    for module in ("kafka", "aiokafka", "confluent_kafka"):
-        _patch_module(module)
+    # Schedule patching of pre-imported modules using a deferred approach
+    # This avoids circular import issues by running after the current import completes
+    import threading
+    def patch_preimported():
+        """Patch any modules that were imported before superclient"""
+        for module in ("kafka", "aiokafka", "confluent_kafka"):
+            if module in sys.modules:
+                _patch_module(module)
+    
+    # Use a timer with 0 delay to run after current import stack completes
+    threading.Timer(0, patch_preimported).start()
 
     # Start heartbeat
     Heartbeat.ensure() 

superclient/__main__.py

@@ -103,15 +103,15 @@ def init_patch(self, *args, **kwargs):
             if not hasattr(self, "_superstream_patch"):
                 original_send = self.send
 
-                def send_patch(inner, topic, *a, **kw):
+                def send_patch(topic, *a, **kw):
                     """Track topic usage when sending messages."""
                     tr.record_topic(topic)
                     return original_send(topic, *a, **kw)
 
                 self.send = send_patch
                 orig_close = self.close
 
-                def close_patch(inner, *a, **kw):
+                def close_patch(*a, **kw):
                     """Clean up Superstream resources when closing the producer."""
                     if not hasattr(self, "_superstream_closed"):
                         self._superstream_closed = True
@@ -211,14 +211,14 @@ def init_patch(self, *args, **kwargs):
             if not hasattr(self, "_superstream_patch"):
                 original_send = self.send
 
-                async def send_patch(inner, topic, *a, **kw):
+                async def send_patch(topic, *a, **kw):
                     tr.record_topic(topic)
                     return await original_send(topic, *a, **kw)
 
                 self.send = send_patch
                 original_stop = self.stop
 
-                async def stop_patch(inner, *a, **kw):
+                async def stop_patch(*a, **kw):
                     if not hasattr(self, "_superstream_closed"):
                         self._superstream_closed = True
                         tr.close()
@@ -310,14 +310,14 @@ def init_patch(self, conf: Dict[str, Any], *args, **kwargs):
             if not hasattr(self, "_superstream_patch"):
                 original_produce = self.produce
 
-                def produce_patch(inner, topic, *a, **kw):
+                def produce_patch(topic, *a, **kw):
                     tr.record_topic(topic)
                     return original_produce(topic, *a, **kw)
 
                 self.produce = produce_patch
                 orig_close = self.close
 
-                def close_patch(inner, *a, **kw):
+                def close_patch(*a, **kw):
                     if not hasattr(self, "_superstream_closed"):
                         self._superstream_closed = True
                         tr.close()

superclient/agent/__init__.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Install the superclient-init.pth file to enable automatic loading"""
+import os
+import sys
+import shutil
+import site
+
+def install_pth():
+    """Install the .pth file to the appropriate site-packages directory"""
+    # Get the path to the .pth file from package resources
+    import superclient
+    package_dir = os.path.dirname(superclient.__file__)
+    source = os.path.join(package_dir, "superclient-init.pth")
+    
+    if not os.path.exists(source):
+        print(f"Error: {source} not found")
+        return False
+    
+    # Find site-packages directory
+    site_packages = None
+    for path in sys.path:
+        if path.endswith('site-packages') and os.path.isdir(path):
+            site_packages = path
+            break
+    
+    if not site_packages:
+        site_packages = site.getsitepackages()[0]
+    
+    dest = os.path.join(site_packages, "superclient-init.pth")
+    
+    try:
+        shutil.copy2(source, dest)
+        print(f"Successfully installed {dest}")
+        print("Superclient will now load automatically!")
+        return True
+    except Exception as e:
+        print(f"Error installing .pth file: {e}")
+        print(f"You may need to run this script with sudo or as administrator")
+        return False
+
+if __name__ == "__main__":
+    install_pth() 

superclient/agent/interceptor.py

@@ -0,0 +1 @@
+import superclient