Provide patch_all and patch as main way of monkey patching

Plus quick update of the doc

ddtrace/__init__.py

@@ -1,9 +1,19 @@
 
+from .monkey import patch, patch_all
+from .pin import Pin
+from .span import Span
 from .tracer import Tracer
-from .span import Span # noqa
-from .pin import Pin   # noqa
 
 __version__ = '0.3.16'
 
-# a global tracer
+# a global tracer instance
 tracer = Tracer()
+
+__all__ = [
+    'patch',
+    'patch_all',
+    'Pin',
+    'Span',
+    'tracer',
+    'Tracer',
+]

ddtrace/monkey.py

@@ -1,8 +1,102 @@
+"""Patch librairies to be automatically instrumented.
 
-from ddtrace.contrib import autopatch
+It can monkey patch supported standard libraries and third party modules.
+A patched module will automatically report spans with its default configuration.
 
-def patch_all():
-    autopatch.autopatch()
+A library instrumentation can be configured (for instance, to report as another service)
+using Pin. For that, check its documentation.
+"""
+import logging
+import importlib
+import threading
+
+
+# Default set of modules to automatically patch or not
+PATCH_MODULES = {
+    'cassandra': True,
+    'elasticsearch': True,
+    'mongoengine': True,
+    'psycopg': True,
+    'pylibmc': True,
+    'pymongo': True,
+    'redis': True,
+    'requests': False,  # Not ready yet
+    'sqlalchemy': False,  # Prefer DB client instrumentation
+    'sqlite3': True,
+}
+
+_LOCK = threading.Lock()
+_PATCHED_MODULES = set()
+
+
+def patch_all(**patch_modules):
+    """Patch all possible modules.
+
+    The list of modules to instrument comes from `PATCH_MODULES`, which
+    is then overridden by `patch_modules`.
+    Calling it multiple times can add more patches, but won't remove
+    existing patches.
+
+    :param dict **patch_modules: override which modules to load or not.
+        Example: {'redis': False, 'cassandra': False}
+    """
+    modules = PATCH_MODULES.copy()
+    modules.update(patch_modules)
+
+    patch(raise_errors=False, **modules)
+
+def patch(raise_errors=True, **patch_modules):
+    """Patch a set of given modules
+
+    :param bool raise_errors: Raise error if one patch fail.
+    :param dict **patch_modules: List of modules to patch.
+        Example: {'psycopg': True, 'elasticsearch': True}
+    """
+    modules = [m for (m, should_patch) in patch_modules.items() if should_patch]
+    count = 0
+    for module in modules:
+        patched = patch_module(module, raise_errors=raise_errors)
+        if patched:
+            count += 1
+
+    logging.info("patched %s/%s modules (%s)",
+        count,
+        len(modules),
+        ",".join(get_patched_modules()))
+
+
+def patch_module(module, raise_errors=True):
+    """Patch a single module
+
+    Returns if the module got properly patched.
+    """
+    try:
+        return _patch_module(module)
+    except Exception as exc:
+        if raise_errors:
+            raise
+        logging.debug("failed to patch %s: %s", module, exc)
+        return False
 
 def get_patched_modules():
-    return autopatch.get_patched_modules()
+    """Get the list of patched modules"""
+    with _LOCK:
+        return sorted(_PATCHED_MODULES)
+
+def _patch_module(module):
+    """_patch_module will attempt to monkey patch the module.
+
+    Returns if the module got patched.
+    Can also raise errors if it fails.
+    """
+    path = 'ddtrace.contrib.%s.patch' % module
+    with _LOCK:
+        if module in _PATCHED_MODULES:
+            logging.debug("already patched: %s", path)
+            return False
+
+        imported_module = importlib.import_module(path)
+        imported_module.patch()
+
+        _PATCHED_MODULES.add(module)
+        return True

docs/index.rst

@@ -30,11 +30,22 @@ We support many `Web Frameworks`_. Install the middleware for yours.
 Then let's patch all the widely used Python libraries that you are running::
 
     # Add the following a the main entry point of your application.
-    from ddtrace import monkey
-    monkey.patch_all()
+    from ddtrace import patch_all
+    patch_all()
 
 Start your web server and you should be off to the races.
 
+If you want to restrict the set of instrumented libraries, you can either say
+which ones to instrument, or which ones not to.
+
+    from ddtrace import patch_all, patch
+
+    # Patch all libraries, except mysql and pymongo
+    patch_all(mysql=False, pymongo=False)
+
+    # Only patch redis and elasticsearch, raising an exception if one fails
+    patch(redis=True, elasticsearch=True, raise_errors=True)
+
 Custom Tracing
 ~~~~~~~~~~~~~~
 

tests/contrib/psycopg/test_psycopg.py

@@ -7,13 +7,11 @@
 from nose.tools import eq_
 
 # project
-from ddtrace import Tracer
-from ddtrace.contrib.psycopg import connection_factory
+from ddtrace.contrib.psycopg import patch_conn, connection_factory
 
 # testing
 from tests.contrib.config import POSTGRES_CONFIG
 from tests.test_tracer import get_dummy_tracer
-from ddtrace.contrib.psycopg import patch_conn
 
 
 TEST_PORT = str(POSTGRES_CONFIG['port'])

tests/autopatch.py → tests/monkey.py

@@ -1,16 +1,22 @@
 """ auto patch things. """
 
-# manual test for autopatching
+# manual test for monkey patching
 import logging
 import sys
 
 # project
 import ddtrace
-from ddtrace.contrib.autopatch import autopatch
 
 # allow logging
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 
 ddtrace.tracer.debug_logging = True
 
-autopatch()
+# Patch nothing
+ddtrace.patch()
+
+# Patch all except Redis
+ddtrace.patch_all(redis=False)
+
+# Patch Redis
+ddtrace.patch(redis=True)