lxml Cleaner with attribute whitelist

attrcleaner.py

#!/usr/bin/env python

from lxml.html.clean import Cleaner
from lxml.html import defs

class AttrWhitelistCleaner(Cleaner):
    """An HTML Cleaner that can use an attribute whitelist.  Defaults to using
    the attributes that are whitelisted by default with ``safe_attrs_only``
    turned on."""
    def __init__(self, **kw):
        self.attr_whitelist = kw.pop('attr_whitelist', set(defs.safe_attrs))
        super(AttrWhitelistCleaner, self).__init__(**kw)

    def __call__(self, doc):
        self.safe_attrs_only = False
        super(AttrWhitelistCleaner, self).__call__(doc)
        if hasattr(doc, 'getroot'):
            doc = doc.getroot()

        whitelist = self.attr_whitelist
        for el in doc.iter():
            attrib = el.attrib
            for aname in attrib.keys():
                if aname not in whitelist:
                    del attrib[aname]

example_whitelist = set(defs.safe_attrs) | set(['flashvars'])

example_cleaner = AttrWhitelistCleaner(attr_whitelist=example_whitelist)