##############################################################################
#
# Copyright (c) 2001, 2002 Zope Corporation and Contributors.
# All Rights Reserved.
#
# This software is subject to the provisions of the Zope Public License,
# Version 2.0 (ZPL).  A copy of the ZPL should accompany this distribution.
# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
# FOR A PARTICULAR PURPOSE
#
##############################################################################
"""Provide a function that can find object references in pickles
"""
from cPickle import Unpickler
from cStringIO import StringIO

def referencesf(p, rootl=None):
    if rootl is None:
        rootl = []
    u = Unpickler(StringIO(p))
    l = len(rootl)
    # If persistent_load is set to a list, the pickler will append
    # persistent ids to it instead of calling it.
    u.persistent_load = rootl
    u.noload()
    try:
        u.noload()
    except:
        # Hm.  We failed to do second load.  Maybe there wasn't a
        # second pickle.  Let's check:
        f = StringIO(p)
        u = Unpickler(f)
        u.persistent_load = []
        u.noload()
        if len(p) > f.tell():
            raise ValueError, 'Error unpickling %r' % p

    # References may have class info, so we need to
    # check for wrapped references.
    for i in range(l, len(rootl)):
        v = rootl[i]
        if v:
            if isinstance(v, tuple):
                v = v[0]
            rootl[i] = v

    return rootl
