lib/portage/util/shelve.py

# Copyright 2020-2021 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

import logging
import pickle
import shelve

logger = logging.getLogger(__name__)


def open_shelve(db_file, flag="r"):
    """
    The optional flag parameter has the same interpretation as the flag
    parameter of dbm.open()
    """
    try:
        db = shelve.open(db_file, flag=flag)
    except ImportError as e:
        # ImportError has different attributes for python2 vs. python3
        if getattr(e, "name", None) == "bsddb" or getattr(e, "message", None) in (
            "No module named bsddb",
            "No module named _bsddb",
        ):
            from bsddb3 import dbshelve

            db = dbshelve.open(db_file)
        else:
            raise

    return db


def dump(args):
    src = open_shelve(args.src, flag="r")
    try:
        with open(args.dest, "wb") as dest:
            for key in src:
                try:
                    value = src[key]
                except KeyError:
                    logger.exception(key)
                    continue
                pickle.dump((key, value), dest)
    finally:
        src.close()


def restore(args):
    dest = open_shelve(args.dest, flag="c")
    try:
        with open(args.src, "rb") as src:
            while True:
                try:
                    k, v = pickle.load(src)
                except EOFError:
                    break
                else:
                    dest[k] = v
    finally:
        dest.close()