def transform_keys(transform, d):
    """
    Transforms the keys in a dict.

    :param transform: If method, calls with key and value, returns new key
                      If dict, maps keys to key values for new key
                      If list, only returns dict with specified keys
                      Else returns original dict
    :param d: dictionary on which we're transforming keys, or list of dicts

    :return: Dictionary with transformed keys

    >>> transform_keys(lambda k, v: k.replace('o', '0'), dict(one=1, two=2))
    {'tw0': 2, '0ne': 1}
    >>> transform_keys({'one': 1}, dict(one=1, two=2))
    {1: 1, 'two': 2}
    >>> transform_keys(['one'], dict(one=1, two=2))
    {'one': 1}
    >>> transform_keys(None, dict(one=1, two=2))
    {'two': 2, 'one': 1}
    """
    if isinstance(d, list):
        return [transform_keys(transform, i) for i in d]

    if callable(transform):
        return {transform(k, v): v for k, v in d.iteritems()}
    elif isinstance(transform, dict):
        return {transform.get(k, k): v for k, v in d.iteritems()}
    elif isinstance(transform, list):
        return {k: v for k, v in d.iteritems() if k in transform}
    else:
        return d