<commit_before>#!/usr/bin/env python3

"""nndb-import.py
"""

"""
These are the files that we support as documented in the file sr27_doc.pdf
as given in the SR27 distribution of these files. The four files identified
as the four "principal" files are the first four listed:

   Name      Descrip
   ========= ==========================================
   FOOD_DES  food descriptions
   NUT_DATA  nutrient data
   WEIGHT    weights
   FOOTNOTE  footnotes
   --------
   FD_GROUP  food group descriptions
   LANGUAL   LanguaL Factors
   LANGDESC  LanguaL Factors descriptions
   NUTR_DEF  Nutrient definitions
   SRC_CD    source code
   DERIV_CD  data derivation code description
   DATA_SRC  Sources of data
   DATSRCLN  Sources of data Link
"""

import sys
import os
import argparse
import collections
import functools

try:
    import pymongo
except:
    sys.stderr.write("\n\nCould not import PyMongo - it is required\n\n")
    raise

def nndb_recs(filename, field_names=None):
    """Iterator that yields a dictionary of (field-name,value) pairs. Note that
    the order of field_names is important!
    """

    if not field_names:
        raise ValueError("No fields specified for nndb recs")

    # Filter for a single field in the file
    def filt(f):
        if len(f) > 1 and f[0] == '~' and f[-1] == '~':
            f = f[1:-1]
        return f

    # Note the encoding - the documentation says the files are in
    # "ASCII (ISO/IEC 8859-1)" which might make you think you could assume
    # an encoding of ASCII or UTF-8, but they actually do use 8859-1 (latin_1)
    # characters that are not valid ASCII or UTF-8.
    with open(filename, "r", encoding="latin_1") as datafile:
        for line in datafile:
            line = line.strip()
            if len(line) < 1:
                continue
            yield dict([
                (name, filt(fld))
                for name, fld in zip(field_names, line.split('^'))
            ])

def recs_to_dict(keyfldname, recs):
    """Given an iterable of recs and a keyfldname, build and return a
    dictionary of records"""
    d = dict()
    for rec in recs:
        d[rec[keyfldname]] = rec
    return d

def recs_to_lookup(filename):
    """Unlike recs_to_dict, we turn a 2-column data file (where the cols are
    assumed to be key and value) into a dictionary
    """
    d = {"": ""}
    for flds in nndb_recs(filename, ["key", "val"]):
        d[flds["key"]] = flds["val"]
    return d

food_des_recs = functools.partial(nndb_recs, field_names=[
    "_id",  # aka NDB_No
    "food_group_code",  # references the food group descriptions
    "descrip",
    "short_descrip",
    "common_name",
    "mfg_name",
    "survey",  # if used in FNDDS (if so, nutrient data should be complete)
    "refuse_descrip",  # description of inedible parts (seed, bone, etc)
    "refuse",  # percentage of refuse
    "scientific_name",  # generally for least processed, raw form if applicable
    "n_factor",  # factor for nitrogen to protein
    "protein_factor",  # factor for calories from protein
    "fat_factor",  # factor for calories from fat
    "carb_facto",  # factor for calories from carbohydrates
])

nut_data_recs = functools.partial(nndb_recs, field_names=[
    "ndb_num",  # aka NDB_No, the _id to FOOD_DES
    "nutrient_id",  # aka Nutr_No
    "nutrient_val",  # Num edible portion in 100g
    "data_point_count",  # Num data points used for analysis (0 means calc, etc)
    "std_error",  # Num - std err of mean, can be null (if data_point_count < 3)
    "source_code",
    "derivation_code",
    "ref_nbd_id",  # May refer to another item used to calc a missing value
    "add_nutrition_mark",  # Used for fortified cereals
    "num_studies",  # Num
    "min_value",  # Num
    "max_value",  # Num
    "degrees_freedom",  # Num
    "lower_err_bound",  # Num - lower 95% error bound
    "upper_err_bound",  # Num - lower 95% error bound
    "statistical_comments",
    "confidence_code",  # Indicates overall assessment of quality for sampling, etc
])

nutr_def_recs = functools.partial(nndb_recs, field_names=[
    "nutrient_id",  # aka Nutr_No
    "units",  # mg, g, lb, etc
    "tagname",  # INFOODS tag
    "descrip",
    "decimal_places",
    "sr_sort_order",  # num
])

def num(s, filt=float):
    """Helper for numeric fields - we accept a string, convert it using filt,
    and will return an empty string if there is a ValueError converting
    """
    if not s:
        return ""
    try:
        return filt(s)
    except ValueError:
        return ""

def nums(rec, field_names, filt=float):
    """Mutate rec using num (with filt) for all given field names
    """
    for fn in field_names:
        rec[fn] = num(rec.get(fn, ""), filt=filt)

def process_directory(mongo, dirname):
    # Shorten our code a little
    get_fn = lambda fn: os.path.join(dirname, fn)

    # Read in various dictionaries we need first
    print("Reading food group codes")
    food_grp_codes = recs_to_lookup(get_fn("FD_GROUP.txt"))

    print("Reading LanguaL codes")
    langual_codes = recs_to_dict(
        "code",
        nndb_recs(get_fn("LANGDESC.txt"), ["code", "descrip"])
    )
    langual_codes[""] = ""

    # First create all the entries with default values from the main file
    # Note that we use upsert - one of main goals is to be restartable and
    # re-runnable.
    print("Creating entries...")
    count = 0
    survey_stats = collections.defaultdict(int, [("", 0), ("Y", 0)])
    for entry in food_des_recs(get_fn("FOOD_DES.txt")):
        # Just in case we ever decide to use something else for _id
        entry['ndb_num'] = entry['_id']
        # Handle numeric fields
        nums(entry, ["n_factor", "protein_factor", "fat_factor", "carb_factor"])
        # Setup defaults
        entry.update({
            'nutrients': list(),
            'food_group_descrip': food_grp_codes[entry["food_group_code"]],
            'footnotes': list(),
            'langual_entries': list(),
            'measures': list(),
        })

        # Upsert our entry (already have an _id, so this will be an upsert)
        mongo.save(entry)

        count += 1
        survey_stats[entry['survey']] += 1
        if count % 3000 == 0:
            print("  Created %7d" % count)

    print("...Total Created: %d" % count)
    print("Survey stats:")
    for k, v in survey_stats.items():
        print("  %4s: %12d" % (k, v))
    print("")

    print("Loading weights/measures")
    weight_recs = nndb_recs(get_fn("WEIGHT.txt"), [
        "ndb_num",
        "seq",
        "amount",
        "descrip",
        "gram_weight",
        "num_data_points",
        "stddev"
    ])
    count = 0
    for entry in weight_recs:
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"measures": entry}}
        )
        count += 1
        if count % 10000 == 0:
            print("  weights: %7d" % count)
    print("...Total weights read: %d" % count)

    print("Loading LanguaL Codes...")
    count = 0
    for entry in nndb_recs(get_fn("LANGUAL.txt"), ["ndb_num", "code"]):
        lang = langual_codes[entry["code"]]
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"langual_entries": lang}}
        )
        count += 1
        if count % 10000 == 0:
            print("  LanguaL items: %7d" % count)
    print("...Total codes read: %d" % count)

    print("Reading source code descrips")
    src_codes = recs_to_lookup(get_fn("SRC_CD.txt"))

    print("Reading data derivation code descrips")
    deriv_codes = recs_to_lookup(get_fn("DERIV_CD.txt"))

    print("Reading data sources for footnotes")
    data_srcs = recs_to_dict("datasrc_id", nndb_recs(get_fn("DATA_SRC.txt"), [
        "datasrc_id",
        "authors",
        "title",
        "year",
        "journal",
        "vol_city",
        "issue_state",
        "start_page",
        "end_page",
    ]))

    print("Adding footnotes to food items...")
    footnote_recs = nndb_recs(get_fn("FOOTNOTE.txt"), [
        "ndb_num",
        "footnote_num",
        "footnote_type",
        "nutr_num",
        "text"
    ])
    count = 0
    for entry in footnote_recs:
        entry["type"] = "footnote"
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"footnotes": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  footnotes: %7d" % count)
    print("...Total foot notes: %7d" % count)

    print("Adding data sources to food items as footnotes...")
    count = 0
    for entry in nndb_recs(get_fn("DATSRCLN.txt"), ["ndb_num", "nutr_num", "datasrc_id"]):
        entry.update(data_srcs[entry["datasrc_id"]])
        entry["type"] = "data-source"
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"footnotes": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  data sources: %7d" % count)
    print("...Total data sources: %7d" % count)

    print("Reading nutrient defs...")
    nutr_defs = recs_to_dict(
        "nutrient_id",
        nutr_def_recs(get_fn("NUTR_DEF.txt"))
    )

    print("Loading nutrition items...")
    count = 0
    for entry in nut_data_recs(get_fn("NUT_DATA.txt")):
        nums(entry, [
            "nutrient_val",
            "data_point_count",
            "std_error",
            "num_studies",
            "min_value", "max_value",
            "degrees_freedom",
            "lower_err_bound", "upper_err_bound",
        ])

        # Add descriptions for the codes we know
        entry["source_descrip"] = src_codes[entry["source_code"]]
        entry["derivation_descrip"] = deriv_codes[entry["derivation_code"]]

        # Just add the nutrient def info to the entry
        xtra = nutr_defs[entry["nutrient_id"]]
        entry.update(xtra)

        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"nutrients": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  nutrient items: %7d" % count)
    print("...Total nutrient items: %7d" % count)

# TODO: we're going to want a check that all nutrients in an entry have their
#       seq num equal to their array index in nutrients

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "targetdir",
        help="The directory containing the ASCII data files from "
    )
    parser.add_argument(
        "--mongourl",
        help="The MongoDB URL to use for connection",
        default="mongodb://localhost:27017/nutrition"
    )
    parser.add_argument(
        "--collection",
        help="The collection where the data will be loaded",
        default="nndb"
    )
    args = parser.parse_args()

    print(args.targetdir)

    print("Connecting to %s" % args.mongourl)
    client = pymongo.MongoClient(args.mongourl)
    db = client.get_default_database()
    print("Using collection %s" % args.collection)
    coll = db[args.collection]
    print("Processing files using directory %s" % args.targetdir)
    process_directory(coll, args.targetdir)

if __name__ == "__main__":
    main()
<commit_msg>todo update<commit_after>#!/usr/bin/env python3

# TODO: translate microgram abbreviation "Âµg" to something easier to type

# TODO: we're going to want a check that all nutrients in an entry have their
#       seq num equal to their array index in nutrients

"""nndb-import.py

These are the files that we support as documented in the file sr27_doc.pdf
as given in the SR27 distribution of these files. The four files identified
as the four "principal" files are the first four listed:

   Name      Descrip
   ========= ==========================================
   FOOD_DES  food descriptions
   NUT_DATA  nutrient data
   WEIGHT    weights
   FOOTNOTE  footnotes
   --------
   FD_GROUP  food group descriptions
   LANGUAL   LanguaL Factors
   LANGDESC  LanguaL Factors descriptions
   NUTR_DEF  Nutrient definitions
   SRC_CD    source code
   DERIV_CD  data derivation code description
   DATA_SRC  Sources of data
   DATSRCLN  Sources of data Link
"""

import sys
import os
import argparse
import collections
import functools

try:
    import pymongo
except:
    sys.stderr.write("\n\nCould not import PyMongo - it is required\n\n")
    raise

def nndb_recs(filename, field_names=None):
    """Iterator that yields a dictionary of (field-name,value) pairs. Note that
    the order of field_names is important!
    """

    if not field_names:
        raise ValueError("No fields specified for nndb recs")

    # Filter for a single field in the file
    def filt(f):
        if len(f) > 1 and f[0] == '~' and f[-1] == '~':
            f = f[1:-1]
        return f

    # Note the encoding - the documentation says the files are in
    # "ASCII (ISO/IEC 8859-1)" which might make you think you could assume
    # an encoding of ASCII or UTF-8, but they actually do use 8859-1 (latin_1)
    # characters that are not valid ASCII or UTF-8.
    with open(filename, "r", encoding="latin_1") as datafile:
        for line in datafile:
            line = line.strip()
            if len(line) < 1:
                continue
            yield dict([
                (name, filt(fld))
                for name, fld in zip(field_names, line.split('^'))
            ])

def recs_to_dict(keyfldname, recs):
    """Given an iterable of recs and a keyfldname, build and return a
    dictionary of records"""
    d = dict()
    for rec in recs:
        d[rec[keyfldname]] = rec
    return d

def recs_to_lookup(filename):
    """Unlike recs_to_dict, we turn a 2-column data file (where the cols are
    assumed to be key and value) into a dictionary
    """
    d = {"": ""}
    for flds in nndb_recs(filename, ["key", "val"]):
        d[flds["key"]] = flds["val"]
    return d

food_des_recs = functools.partial(nndb_recs, field_names=[
    "_id",  # aka NDB_No
    "food_group_code",  # references the food group descriptions
    "descrip",
    "short_descrip",
    "common_name",
    "mfg_name",
    "survey",  # if used in FNDDS (if so, nutrient data should be complete)
    "refuse_descrip",  # description of inedible parts (seed, bone, etc)
    "refuse",  # percentage of refuse
    "scientific_name",  # generally for least processed, raw form if applicable
    "n_factor",  # factor for nitrogen to protein
    "protein_factor",  # factor for calories from protein
    "fat_factor",  # factor for calories from fat
    "carb_facto",  # factor for calories from carbohydrates
])

nut_data_recs = functools.partial(nndb_recs, field_names=[
    "ndb_num",  # aka NDB_No, the _id to FOOD_DES
    "nutrient_id",  # aka Nutr_No
    "nutrient_val",  # Num edible portion in 100g
    "data_point_count",  # Num data points used for analysis (0 means calc, etc)
    "std_error",  # Num - std err of mean, can be null (if data_point_count < 3)
    "source_code",
    "derivation_code",
    "ref_nbd_id",  # May refer to another item used to calc a missing value
    "add_nutrition_mark",  # Used for fortified cereals
    "num_studies",  # Num
    "min_value",  # Num
    "max_value",  # Num
    "degrees_freedom",  # Num
    "lower_err_bound",  # Num - lower 95% error bound
    "upper_err_bound",  # Num - lower 95% error bound
    "statistical_comments",
    "confidence_code",  # Indicates overall assessment of quality for sampling, etc
])

nutr_def_recs = functools.partial(nndb_recs, field_names=[
    "nutrient_id",  # aka Nutr_No
    "units",  # mg, g, lb, etc
    "tagname",  # INFOODS tag
    "descrip",
    "decimal_places",
    "sr_sort_order",  # num
])

def num(s, filt=float):
    """Helper for numeric fields - we accept a string, convert it using filt,
    and will return an empty string if there is a ValueError converting
    """
    if not s:
        return ""
    try:
        return filt(s)
    except ValueError:
        return ""

def nums(rec, field_names, filt=float):
    """Mutate rec using num (with filt) for all given field names
    """
    for fn in field_names:
        rec[fn] = num(rec.get(fn, ""), filt=filt)

def process_directory(mongo, dirname):
    # Shorten our code a little
    get_fn = lambda fn: os.path.join(dirname, fn)

    # Read in various dictionaries we need first
    print("Reading food group codes")
    food_grp_codes = recs_to_lookup(get_fn("FD_GROUP.txt"))

    print("Reading LanguaL codes")
    langual_codes = recs_to_dict(
        "code",
        nndb_recs(get_fn("LANGDESC.txt"), ["code", "descrip"])
    )
    langual_codes[""] = ""

    # First create all the entries with default values from the main file
    # Note that we use upsert - one of main goals is to be restartable and
    # re-runnable.
    print("Creating entries...")
    count = 0
    survey_stats = collections.defaultdict(int, [("", 0), ("Y", 0)])
    for entry in food_des_recs(get_fn("FOOD_DES.txt")):
        # Just in case we ever decide to use something else for _id
        entry['ndb_num'] = entry['_id']
        # Handle numeric fields
        nums(entry, ["n_factor", "protein_factor", "fat_factor", "carb_factor"])
        # Setup defaults
        entry.update({
            'nutrients': list(),
            'food_group_descrip': food_grp_codes[entry["food_group_code"]],
            'footnotes': list(),
            'langual_entries': list(),
            'measures': list(),
        })

        # Upsert our entry (already have an _id, so this will be an upsert)
        mongo.save(entry)

        count += 1
        survey_stats[entry['survey']] += 1
        if count % 3000 == 0:
            print("  Created %7d" % count)

    print("...Total Created: %d" % count)
    print("Survey stats:")
    for k, v in survey_stats.items():
        print("  %4s: %12d" % (k, v))
    print("")

    print("Loading weights/measures")
    weight_recs = nndb_recs(get_fn("WEIGHT.txt"), [
        "ndb_num",
        "seq",
        "amount",
        "descrip",
        "gram_weight",
        "num_data_points",
        "stddev"
    ])
    count = 0
    for entry in weight_recs:
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"measures": entry}}
        )
        count += 1
        if count % 10000 == 0:
            print("  weights: %7d" % count)
    print("...Total weights read: %d" % count)

    print("Loading LanguaL Codes...")
    count = 0
    for entry in nndb_recs(get_fn("LANGUAL.txt"), ["ndb_num", "code"]):
        lang = langual_codes[entry["code"]]
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"langual_entries": lang}}
        )
        count += 1
        if count % 10000 == 0:
            print("  LanguaL items: %7d" % count)
    print("...Total codes read: %d" % count)

    print("Reading source code descrips")
    src_codes = recs_to_lookup(get_fn("SRC_CD.txt"))

    print("Reading data derivation code descrips")
    deriv_codes = recs_to_lookup(get_fn("DERIV_CD.txt"))

    print("Reading data sources for footnotes")
    data_srcs = recs_to_dict("datasrc_id", nndb_recs(get_fn("DATA_SRC.txt"), [
        "datasrc_id",
        "authors",
        "title",
        "year",
        "journal",
        "vol_city",
        "issue_state",
        "start_page",
        "end_page",
    ]))

    print("Adding footnotes to food items...")
    footnote_recs = nndb_recs(get_fn("FOOTNOTE.txt"), [
        "ndb_num",
        "footnote_num",
        "footnote_type",
        "nutr_num",
        "text"
    ])
    count = 0
    for entry in footnote_recs:
        entry["type"] = "footnote"
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"footnotes": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  footnotes: %7d" % count)
    print("...Total foot notes: %7d" % count)

    print("Adding data sources to food items as footnotes...")
    count = 0
    for entry in nndb_recs(get_fn("DATSRCLN.txt"), ["ndb_num", "nutr_num", "datasrc_id"]):
        entry.update(data_srcs[entry["datasrc_id"]])
        entry["type"] = "data-source"
        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"footnotes": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  data sources: %7d" % count)
    print("...Total data sources: %7d" % count)

    print("Reading nutrient defs...")
    nutr_defs = recs_to_dict(
        "nutrient_id",
        nutr_def_recs(get_fn("NUTR_DEF.txt"))
    )

    print("Loading nutrition items...")
    count = 0
    for entry in nut_data_recs(get_fn("NUT_DATA.txt")):
        nums(entry, [
            "nutrient_val",
            "data_point_count",
            "std_error",
            "num_studies",
            "min_value", "max_value",
            "degrees_freedom",
            "lower_err_bound", "upper_err_bound",
        ])

        # Add descriptions for the codes we know
        entry["source_descrip"] = src_codes[entry["source_code"]]
        entry["derivation_descrip"] = deriv_codes[entry["derivation_code"]]

        # Just add the nutrient def info to the entry
        xtra = nutr_defs[entry["nutrient_id"]]
        entry.update(xtra)

        mongo.update(
            {"_id": entry["ndb_num"]},
            {"$push": {"nutrients": entry}}
        )
        count += 1
        if count % 50000 == 0:
            print("  nutrient items: %7d" % count)
    print("...Total nutrient items: %7d" % count)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "targetdir",
        help="The directory containing the ASCII data files from "
    )
    parser.add_argument(
        "--mongourl",
        help="The MongoDB URL to use for connection",
        default="mongodb://localhost:27017/nutrition"
    )
    parser.add_argument(
        "--collection",
        help="The collection where the data will be loaded",
        default="nndb"
    )
    args = parser.parse_args()

    print(args.targetdir)

    print("Connecting to %s" % args.mongourl)
    client = pymongo.MongoClient(args.mongourl)
    db = client.get_default_database()
    print("Using collection %s" % args.collection)
    coll = db[args.collection]
    print("Processing files using directory %s" % args.targetdir)
    process_directory(coll, args.targetdir)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# © 2017 Houssine BAKKALI, Coop IT Easy
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
from datetime import datetime

from openerp import api, fields, models, _
import openerp.addons.decimal_precision as dp

class product_product(models.Model):
    _inherit = 'product.template'

    scale_group_id = fields.Many2one(related="product_variant_ids.scale_group_id", string='Scale Group', store=True)
    scale_sequence = fields.Integer(related="product_variant_ids.scale_sequence",string='Scale Sequence', store=True)
    scale_tare_weight = fields.Float(related="product_variant_ids.scale_tare_weight",string='Scale Tare Weight', store=True,
        help="Set here Constant tare weight"
        " for the given product. This tare will be substracted when"
        " the product is weighted. Usefull only for weightable product.\n"
        "The tare is defined with kg uom.")

    # View Section
    @api.multi
    def send_scale_create(self):
        for product in self:
            self._send_to_scale_bizerba('create', product)
        return True

    @api.multi
    def send_scale_write(self):
        for product in self:
            self._send_to_scale_bizerba('write', product)
        return True

    @api.multi
    def send_scale_unlink(self):
        for product in self:
            self._send_to_scale_bizerba('unlink', product)
        return True

    # Custom Section
    def _send_to_scale_bizerba(self, action, product):
        log_obj = self.env['product.scale.log']
        log_obj.create({
            'log_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'scale_system_id': product.scale_group_id.scale_system_id.id,
            'product_id': product.id,
            'action': action,
            })

    def _check_vals_scale_bizerba(self, vals, product):
        system = product.scale_group_id.scale_system_id
        system_fields = [x.name for x in system.field_ids]
        vals_fields = vals.keys()
        return set(system_fields).intersection(vals_fields)

    # Overload Section
    @api.model
    def create(self, vals):
        send_to_scale = vals.get('scale_group_id', False)
        res = super(product_product, self).create(
            cr, uid, vals, context=context)
        if send_to_scale:
            product = self.browse(res)
            self._send_to_scale_bizerba('create', product)
        return res

    @api.multi
    def write(self):
        defered = {}
        for product in self:
            ignore = not product.scale_group_id\
                and 'scale_group_id' not in vals.keys()
            if not ignore:
                if not product.scale_group_id:
                    # (the product is new on this group)
                    defered[product.id] = 'create'
                else:
                    if vals.get('scale_group_id', False) and (
                            vals.get('scale_group_id', False) !=
                            product.scale_group_id):
                        # (the product has moved from a group to another)
                        # Remove from obsolete group
                        self._send_to_scale_bizerba('unlink', product)
                        # Create in the new group
                        defered[product.id] = 'create'
                    elif self._check_vals_scale_bizerba(vals, product):
                        # Data related to the scale
                        defered[product.id] = 'write'

        res = super(product_product, self).write(vals)

        for product_id, action in defered.iteritems():
            product = self.browse(product_id)
            self._send_to_scale_bizerba(action, product)

        return res

    @api.multi
    def unlink(self):
        for product in self:
            if product.scale_group_id:
                self._send_to_scale_bizerba(
                    cr, uid, 'unlink', product, context=context)
        return super(product_product, self).unlink()<commit_msg>[FIX] fix wrong function signature<commit_after># -*- coding: utf-8 -*-
# © 2017 Houssine BAKKALI, Coop IT Easy
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
from datetime import datetime

from openerp import api, fields, models, _
import openerp.addons.decimal_precision as dp

class product_product(models.Model):
    _inherit = 'product.template'

    scale_group_id = fields.Many2one(related="product_variant_ids.scale_group_id", string='Scale Group', store=True)
    scale_sequence = fields.Integer(related="product_variant_ids.scale_sequence",string='Scale Sequence', store=True)
    scale_tare_weight = fields.Float(related="product_variant_ids.scale_tare_weight",string='Scale Tare Weight', store=True,
        help="Set here Constant tare weight"
        " for the given product. This tare will be substracted when"
        " the product is weighted. Usefull only for weightable product.\n"
        "The tare is defined with kg uom.")

    # View Section
    @api.multi
    def send_scale_create(self):
        for product in self:
            self._send_to_scale_bizerba('create', product)
        return True

    @api.multi
    def send_scale_write(self):
        for product in self:
            self._send_to_scale_bizerba('write', product)
        return True

    @api.multi
    def send_scale_unlink(self):
        for product in self:
            self._send_to_scale_bizerba('unlink', product)
        return True

    # Custom Section
    def _send_to_scale_bizerba(self, action, product):
        log_obj = self.env['product.scale.log']
        log_obj.create({
            'log_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'scale_system_id': product.scale_group_id.scale_system_id.id,
            'product_id': product.id,
            'action': action,
            })

    def _check_vals_scale_bizerba(self, vals, product):
        system = product.scale_group_id.scale_system_id
        system_fields = [x.name for x in system.field_ids]
        vals_fields = vals.keys()
        return set(system_fields).intersection(vals_fields)

    # Overload Section
    @api.model
    def create(self, vals):
        send_to_scale = vals.get('scale_group_id', False)
        res = super(product_product, self).create(vals)
        if send_to_scale:
            product = self.browse(res)
            self._send_to_scale_bizerba('create', product)
        return res

    @api.multi
    def write(self):
        defered = {}
        for product in self:
            ignore = not product.scale_group_id\
                and 'scale_group_id' not in vals.keys()
            if not ignore:
                if not product.scale_group_id:
                    # (the product is new on this group)
                    defered[product.id] = 'create'
                else:
                    if vals.get('scale_group_id', False) and (
                            vals.get('scale_group_id', False) !=
                            product.scale_group_id):
                        # (the product has moved from a group to another)
                        # Remove from obsolete group
                        self._send_to_scale_bizerba('unlink', product)
                        # Create in the new group
                        defered[product.id] = 'create'
                    elif self._check_vals_scale_bizerba(vals, product):
                        # Data related to the scale
                        defered[product.id] = 'write'

        res = super(product_product, self).write(vals)

        for product_id, action in defered.iteritems():
            product = self.browse(product_id)
            self._send_to_scale_bizerba(action, product)

        return res

    @api.multi
    def unlink(self):
        for product in self:
            if product.scale_group_id:
                self._send_to_scale_bizerba('unlink', product)
        return super(product_product, self).unlink()<|endoftext|>


<commit_before>import os
import csv
from unidecode import unidecode

from optparse import make_option

from django.core.management.base import (
    BaseCommand,
    CommandError,
)

from apps.api.models import (
    GroupMember,
    Contest,
    Singer,
    Quartet,
)

class Command(BaseCommand):
    help = "Command to import contestants"
    option_list = BaseCommand.option_list + (
        make_option(
            "-f",
            "--file",
            dest="filename",
            help="specify import file",
            metavar="FILE"
        ),
    )

    def c(self, input):
        return unidecode(input.strip())

    def handle(self, *args, **options):
        # make sure file option is present
        if options['filename'] is None:
            raise CommandError("Option `--file=...` must be specified.")

        # make sure file path resolves
        if not os.path.isfile(options['filename']):
            raise CommandError("File does not exist at the specified path.")

        # print file
        print "Path: `%s`" % options['filename']

        # open the file
        with open(options['filename']) as csv_file:
            reader = csv.reader(csv_file)
            next(reader)
            for row in reader:
                try:
                    contest = Contest.objects.get(pk=row[2])
                    quartet = Quartet.objects.get(pk=row[3])
                    singer = Singer.objects.get(pk=row[4])
                    gm = GroupMember.objects.create(
                        part=row[1],
                        contest=contest,
                        quartet=quartet,
                        singer=singer,
                    )
                    print gm
                except Exception as e:
                    print "GM {0} could not be created.".format(
                        row[0],
                    )
                    print "Exception: {0} ".format(e)
        return "Done"
<commit_msg>Fix import script<commit_after>import os
import csv
from unidecode import unidecode

from optparse import make_option

from django.core.management.base import (
    BaseCommand,
    CommandError,
)

from apps.api.models import (
    GroupMember,
    Contest,
    Singer,
    Quartet,
)

class Command(BaseCommand):
    help = "Command to import contestants"
    option_list = BaseCommand.option_list + (
        make_option(
            "-f",
            "--file",
            dest="filename",
            help="specify import file",
            metavar="FILE"
        ),
    )

    def c(self, input):
        return unidecode(input.strip())

    def handle(self, *args, **options):
        # make sure file option is present
        if options['filename'] is None:
            raise CommandError("Option `--file=...` must be specified.")

        # make sure file path resolves
        if not os.path.isfile(options['filename']):
            raise CommandError("File does not exist at the specified path.")

        # print file
        print "Path: `%s`" % options['filename']

        # open the file
        with open(options['filename']) as csv_file:
            reader = csv.reader(csv_file)
            next(reader)
            for row in reader:
                try:
                    contest = Contest.objects.get(pk=row[2])
                    quartet = Quartet.objects.get(pk=row[3])
                    singer = Singer.objects.get(pk=row[4])
                    gm = GroupMember.objects.create(
                        part=row[1],
                        contest=contest,
                        group=quartet,
                        singer=singer,
                    )
                    print gm
                except Exception as e:
                    print "GM {0} could not be created.".format(
                        row[0],
                    )
                    print "Exception: {0} ".format(e)
        return "Done"
<|endoftext|>


<commit_before># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2021, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import logging
import re

from torch import nn
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_MASKED_LM_MAPPING,
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
    TOKENIZER_MAPPING,
    BertConfig,
    BertForMaskedLM,
    BertForSequenceClassification,
    BertForQuestionAnswering,
    BertPreTrainedModel,
    BertTokenizer,
    BertTokenizerFast,
)
from transformers.models.bert.modeling_bert import BertOnlyMLMHead

# This is the dict of the `models` module. Anything imported from `models`
# must be contained in this dict. Hence, it will be updated to include
# the new modules made below so they may be imported.
from . import __dict__ as __models_dict__

# Keep a copy of the mappings so they may be directly accessed.
__models_dict__["CONFIG_MAPPING"] = CONFIG_MAPPING
__models_dict__["MODEL_FOR_MASKED_LM_MAPPING"] = MODEL_FOR_MASKED_LM_MAPPING
__models_dict__["MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING"] = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING # noqa E501
__models_dict__["MODEL_FOR_QUESTION_ANSWERING"] = MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING
__models_dict__["TOKENIZER_MAPPING"] = TOKENIZER_MAPPING

def register_bert_model(bert_cls):
    """
    This function wraps a BertModel inherited cls and automatically:
        1. Creates an associated BertConfig
        2. Creates an associated BertForMaskedLM
        3. Creates an associated BertForSequenceClassification
        4. Registers these classes with Transformers model mappings

    This last step ensures that the resulting config and models may be used by
    AutoConfig, AutoModelForMaskedLM, and AutoModelForSequenceClassification.

    Assumptions are made to auto-name these classes and the corresponding model type.
    For instance, SparseBertModel will have model_type="sparse_bert" and associated
    classes like SparseBertConfig.

    To customize the the inputs to the model's config, include the dataclass
    `bert_cls.ConfigKWargs`. This is, in fact, required. Upon initialization of the
    config, the fields of that dataclass will be used to extract extra keyword arguments
    and assign them as attributes to the config.

    Example
    ```
    @register_bert_model
    class SparseBertModel(BertModel):

        @dataclass
        class ConfigKWargs:
            # Keyword arguments to configure sparsity.
            sparsity: float = 0.9

        # Define __init__, ect.
        ...

    # Model is ready to auto load.
    config = AutoConfig.for_model("sparse_bert", sparsity=0.5)
    model = AutoModelForMaskedLM.from_config(model)

    config.sparsity
    >>> 0.5

    type(model)
    >>> SparseBertModelForMaskedLM
    """

    assert bert_cls.__name__.endswith("BertModel")

    # Get first part of name e.g. StaticSparseBertModel -> StaticSparse
    name_prefix = bert_cls.__name__.replace("BertModel", "")

    # Create new bert config and models based off of `bert_cls`.
    config_cls = create_config_class(bert_cls, name_prefix)
    masked_lm_cls = create_masked_lm_class(bert_cls, name_prefix)
    seq_classification_cls = create_sequence_classification_class(bert_cls, name_prefix)
    # question_answering_cls = create_question_answering_class(bert_cls, name_prefix)

    # Specify the correct config class
    bert_cls.config_class = config_cls
    masked_lm_cls.config_class = config_cls
    seq_classification_cls.config_class = config_cls
    # question_answering_cls.config_class = config_cls

    # Update Transformers mappings to auto-load these new models.
    CONFIG_MAPPING.update({
        config_cls.model_type: config_cls
    })
    TOKENIZER_MAPPING.update({
        config_cls: (BertTokenizer, BertTokenizerFast),
    })
    MODEL_FOR_MASKED_LM_MAPPING.update({
        config_cls: masked_lm_cls,
    })
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.update({
        config_cls: seq_classification_cls
    })
    # MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.update({
    #     config_cls: question_answering_cls
    # })

    # Update the `models` modules so that these classes may be imported.
    __models_dict__.update({
        config_cls.__name__: config_cls,
        masked_lm_cls.__name__: masked_lm_cls,
        seq_classification_cls.__name__: seq_classification_cls,
        # question_answering_cls.__name__: question_answering_cls,
    })

def create_config_class(bert_cls, name_prefix):
    assert hasattr(bert_cls, "ConfigKWargs")
    config_dataclass = bert_cls.ConfigKWargs
    dataclass_fields = config_dataclass.__dataclass_fields__

    class NewBertConfig(BertConfig):
        """
        This is the configuration class to store the configuration of a
        {name_prefix}BertModel. This is adapted from `BertConfig`_.

        .. _BertConfig:
            https://huggingface.co/transformers/_modules/transformers/models/bert/configuration_bert.html#BertConfig
        """

        def __init__(
            self,
            vocab_size=30522,
            hidden_size=768,
            num_hidden_layers=12,
            num_attention_heads=12,
            intermediate_size=3072,
            hidden_act="gelu",
            hidden_dropout_prob=0.1,
            attention_probs_dropout_prob=0.1,
            max_position_embeddings=512,
            type_vocab_size=2,
            initializer_range=0.02,
            layer_norm_eps=1e-12,
            pad_token_id=0,
            gradient_checkpointing=False,
            position_embedding_type="absolute",
            use_cache=True,
            **kwargs
        ):
            super().__init__(pad_token_id=pad_token_id, **kwargs)

            # Bert config params.
            self.vocab_size = vocab_size
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.hidden_act = hidden_act
            self.intermediate_size = intermediate_size
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.initializer_range = initializer_range
            self.layer_norm_eps = layer_norm_eps
            self.gradient_checkpointing = gradient_checkpointing
            self.position_embedding_type = position_embedding_type
            self.use_cache = use_cache

            # Use the dataclass fields to define extra config parameters.
            extra_config = {}
            for arg in dataclass_fields.keys():
                if arg in kwargs:
                    extra_config[arg] = kwargs[arg]

            extra_config = config_dataclass(**extra_config).__dict__
            self.__dict__.update(extra_config)

    # Format model_type e.g. "StaticSparse" will become "static_sparse_bert"
    name_segments = re.findall("[A-Z][^A-Z]*", name_prefix)
    model_type = "_".join(name_segments).lower()
    model_type += "_bert"

    # Rename new config class, format docstring, and add model_type.
    new_cls = NewBertConfig
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)
    new_cls.model_type = model_type

    return new_cls

def create_masked_lm_class(bert_cls, name_prefix):
    """
    Create a BertForMaskedLM that calls `bert_cls` in it's forward.
    """

    class NewBertForMaskedLM(BertForMaskedLM):
        """
        Sparse Bert Model with a `language modeling` head on top. Calls
        {name_prefix}BertModel in its forward.

        This is adapted from `BertForMaskedLM`_

        .. _SparseBertForMaskedLM:
            https://huggingface.co/transformers/_modules/transformers/models/bert/modeling_bert.html#BertForMaskedLM
        """

        def __init__(self, config):

            # Call the init one parent class up.
            # Otherwise, the model will be defined twice.
            BertPreTrainedModel.__init__(self, config)

            if config.is_decoder:
                logging.warning(
                    # This warning was included with the original BertForMaskedLM.
                    f"If you want to use `{name_prefix}BertForMaskedLM` make sure "
                    " `config.is_decoder=False` for bi-directional self-attention."
                )

            self.bert = bert_cls(config, add_pooling_layer=False)
            self.cls = BertOnlyMLMHead(config)

            self.init_weights()

    # Rename new model class and format docstring.
    new_cls = NewBertForMaskedLM
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

    return new_cls

def create_sequence_classification_class(bert_cls, name_prefix):
    """
    Create a BertForSequenceClassification that calls `bert_cls` in it's forward.
    """

    class NewBertForSequenceClassification(BertForSequenceClassification):
        """
        Sparse Bert Model transformer with a sequence classification/regression head on
        top (a linear layer on top of the pooled output) e.g. for GLUE tasks. Calls
        {name_prefix}BertModel in forward.

        This is adapted from `BertForMaskedLM`_

        .. BertForSequenceClassification:
            https://huggingface.co/transformers/_modules/transformers/models/bert/modeling_bert.html#BertForSequenceClassification
        """

        def __init__(self, config):

            # Call the init one parent class up.
            # Otherwise, the model will be defined twice.
            BertPreTrainedModel.__init__(self, config)

            self.num_labels = config.num_labels

            # Replace `BertModel` with SparseBertModel.
            self.bert = bert_cls(config)
            self.dropout = nn.Dropout(config.hidden_dropout_prob)
            self.classifier = nn.Linear(config.hidden_size, config.num_labels)

            self.init_weights()

    # Rename new model class and format docstring.
    new_cls = NewBertForSequenceClassification
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

    return new_cls

# def create_question_answering_class(bert_cls, name_prefix):
#     """
#     Create a BertForQuestionAnswering that calls `bert_cls` in it's forward.
#     """

#     class NewBertForQuestionAnswering(BertForQuestionAnswering):
#         """
#         Bert Model with a span classification head on top for extractive
#         question-answering tasks like SQuAD (a linear layers on top of the
#         hidden-states output to compute `span start logits` and
#         `span end logits`).

#         https://huggingface.co/transformers/model_doc/bert.html#transformers.models.bert.modeling_bert.BertForQuestionAnswering  # noqa: E501
#         """

#         def __init__(self, config):

#             # import pdb
#             # pdb.set_trace()

#             # Call the init one parent class up.
#             # Otherwise, the model will be defined twice.
#             BertPreTrainedModel.__init__(self, config)

#             self.num_labels = config.num_labels

#             # Replace `BertModel` with SparseBertModel.
#             self.bert = bert_cls(config)
#             self.dropout = nn.Dropout(config.hidden_dropout_prob)
#             self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

#             self.init_weights()

#     # Rename new model class and format docstring.
#     new_cls = NewBertForQuestionAnswering
#     new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
#     new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

#     return new_cls<commit_msg>uncomment register bert model stuff<commit_after># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2021, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import logging
import re

from torch import nn
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_MASKED_LM_MAPPING,
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
    TOKENIZER_MAPPING,
    BertConfig,
    BertForMaskedLM,
    BertForSequenceClassification,
    BertForQuestionAnswering,
    BertPreTrainedModel,
    BertTokenizer,
    BertTokenizerFast,
)
from transformers.models.bert.modeling_bert import BertOnlyMLMHead

# This is the dict of the `models` module. Anything imported from `models`
# must be contained in this dict. Hence, it will be updated to include
# the new modules made below so they may be imported.
from . import __dict__ as __models_dict__

# Keep a copy of the mappings so they may be directly accessed.
__models_dict__["CONFIG_MAPPING"] = CONFIG_MAPPING
__models_dict__["MODEL_FOR_MASKED_LM_MAPPING"] = MODEL_FOR_MASKED_LM_MAPPING
__models_dict__["MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING"] = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING # noqa E501
__models_dict__["MODEL_FOR_QUESTION_ANSWERING"] = MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING
__models_dict__["TOKENIZER_MAPPING"] = TOKENIZER_MAPPING

def register_bert_model(bert_cls):
    """
    This function wraps a BertModel inherited cls and automatically:
        1. Creates an associated BertConfig
        2. Creates an associated BertForMaskedLM
        3. Creates an associated BertForSequenceClassification
        4. Registers these classes with Transformers model mappings

    This last step ensures that the resulting config and models may be used by
    AutoConfig, AutoModelForMaskedLM, and AutoModelForSequenceClassification.

    Assumptions are made to auto-name these classes and the corresponding model type.
    For instance, SparseBertModel will have model_type="sparse_bert" and associated
    classes like SparseBertConfig.

    To customize the the inputs to the model's config, include the dataclass
    `bert_cls.ConfigKWargs`. This is, in fact, required. Upon initialization of the
    config, the fields of that dataclass will be used to extract extra keyword arguments
    and assign them as attributes to the config.

    Example
    ```
    @register_bert_model
    class SparseBertModel(BertModel):

        @dataclass
        class ConfigKWargs:
            # Keyword arguments to configure sparsity.
            sparsity: float = 0.9

        # Define __init__, ect.
        ...

    # Model is ready to auto load.
    config = AutoConfig.for_model("sparse_bert", sparsity=0.5)
    model = AutoModelForMaskedLM.from_config(model)

    config.sparsity
    >>> 0.5

    type(model)
    >>> SparseBertModelForMaskedLM
    """

    assert bert_cls.__name__.endswith("BertModel")

    # Get first part of name e.g. StaticSparseBertModel -> StaticSparse
    name_prefix = bert_cls.__name__.replace("BertModel", "")

    # Create new bert config and models based off of `bert_cls`.
    config_cls = create_config_class(bert_cls, name_prefix)
    masked_lm_cls = create_masked_lm_class(bert_cls, name_prefix)
    seq_classification_cls = create_sequence_classification_class(bert_cls, name_prefix)
    question_answering_cls = create_question_answering_class(bert_cls, name_prefix)

    # Specify the correct config class
    bert_cls.config_class = config_cls
    masked_lm_cls.config_class = config_cls
    seq_classification_cls.config_class = config_cls
    question_answering_cls.config_class = config_cls

    # Update Transformers mappings to auto-load these new models.
    CONFIG_MAPPING.update({
        config_cls.model_type: config_cls
    })
    TOKENIZER_MAPPING.update({
        config_cls: (BertTokenizer, BertTokenizerFast),
    })
    MODEL_FOR_MASKED_LM_MAPPING.update({
        config_cls: masked_lm_cls,
    })
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.update({
        config_cls: seq_classification_cls
    })
    MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.update({
        config_cls: question_answering_cls
    })

    # Update the `models` modules so that these classes may be imported.
    __models_dict__.update({
        config_cls.__name__: config_cls,
        masked_lm_cls.__name__: masked_lm_cls,
        seq_classification_cls.__name__: seq_classification_cls,
        question_answering_cls.__name__: question_answering_cls,
    })

def create_config_class(bert_cls, name_prefix):
    assert hasattr(bert_cls, "ConfigKWargs")
    config_dataclass = bert_cls.ConfigKWargs
    dataclass_fields = config_dataclass.__dataclass_fields__

    class NewBertConfig(BertConfig):
        """
        This is the configuration class to store the configuration of a
        {name_prefix}BertModel. This is adapted from `BertConfig`_.

        .. _BertConfig:
            https://huggingface.co/transformers/_modules/transformers/models/bert/configuration_bert.html#BertConfig
        """

        def __init__(
            self,
            vocab_size=30522,
            hidden_size=768,
            num_hidden_layers=12,
            num_attention_heads=12,
            intermediate_size=3072,
            hidden_act="gelu",
            hidden_dropout_prob=0.1,
            attention_probs_dropout_prob=0.1,
            max_position_embeddings=512,
            type_vocab_size=2,
            initializer_range=0.02,
            layer_norm_eps=1e-12,
            pad_token_id=0,
            gradient_checkpointing=False,
            position_embedding_type="absolute",
            use_cache=True,
            **kwargs
        ):
            super().__init__(pad_token_id=pad_token_id, **kwargs)

            # Bert config params.
            self.vocab_size = vocab_size
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.hidden_act = hidden_act
            self.intermediate_size = intermediate_size
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.initializer_range = initializer_range
            self.layer_norm_eps = layer_norm_eps
            self.gradient_checkpointing = gradient_checkpointing
            self.position_embedding_type = position_embedding_type
            self.use_cache = use_cache

            # Use the dataclass fields to define extra config parameters.
            extra_config = {}
            for arg in dataclass_fields.keys():
                if arg in kwargs:
                    extra_config[arg] = kwargs[arg]

            extra_config = config_dataclass(**extra_config).__dict__
            self.__dict__.update(extra_config)

    # Format model_type e.g. "StaticSparse" will become "static_sparse_bert"
    name_segments = re.findall("[A-Z][^A-Z]*", name_prefix)
    model_type = "_".join(name_segments).lower()
    model_type += "_bert"

    # Rename new config class, format docstring, and add model_type.
    new_cls = NewBertConfig
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)
    new_cls.model_type = model_type

    return new_cls

def create_masked_lm_class(bert_cls, name_prefix):
    """
    Create a BertForMaskedLM that calls `bert_cls` in it's forward.
    """

    class NewBertForMaskedLM(BertForMaskedLM):
        """
        Sparse Bert Model with a `language modeling` head on top. Calls
        {name_prefix}BertModel in its forward.

        This is adapted from `BertForMaskedLM`_

        .. _SparseBertForMaskedLM:
            https://huggingface.co/transformers/_modules/transformers/models/bert/modeling_bert.html#BertForMaskedLM
        """

        def __init__(self, config):

            # Call the init one parent class up.
            # Otherwise, the model will be defined twice.
            BertPreTrainedModel.__init__(self, config)

            if config.is_decoder:
                logging.warning(
                    # This warning was included with the original BertForMaskedLM.
                    f"If you want to use `{name_prefix}BertForMaskedLM` make sure "
                    " `config.is_decoder=False` for bi-directional self-attention."
                )

            self.bert = bert_cls(config, add_pooling_layer=False)
            self.cls = BertOnlyMLMHead(config)

            self.init_weights()

    # Rename new model class and format docstring.
    new_cls = NewBertForMaskedLM
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

    return new_cls

def create_sequence_classification_class(bert_cls, name_prefix):
    """
    Create a BertForSequenceClassification that calls `bert_cls` in it's forward.
    """

    class NewBertForSequenceClassification(BertForSequenceClassification):
        """
        Sparse Bert Model transformer with a sequence classification/regression head on
        top (a linear layer on top of the pooled output) e.g. for GLUE tasks. Calls
        {name_prefix}BertModel in forward.

        This is adapted from `BertForMaskedLM`_

        .. BertForSequenceClassification:
            https://huggingface.co/transformers/_modules/transformers/models/bert/modeling_bert.html#BertForSequenceClassification
        """

        def __init__(self, config):

            # Call the init one parent class up.
            # Otherwise, the model will be defined twice.
            BertPreTrainedModel.__init__(self, config)

            self.num_labels = config.num_labels

            # Replace `BertModel` with SparseBertModel.
            self.bert = bert_cls(config)
            self.dropout = nn.Dropout(config.hidden_dropout_prob)
            self.classifier = nn.Linear(config.hidden_size, config.num_labels)

            self.init_weights()

    # Rename new model class and format docstring.
    new_cls = NewBertForSequenceClassification
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

    return new_cls

def create_question_answering_class(bert_cls, name_prefix):
    """
    Create a BertForQuestionAnswering that calls `bert_cls` in it's forward.
    """

    class NewBertForQuestionAnswering(BertForQuestionAnswering):
        """
        Bert Model with a span classification head on top for extractive
        question-answering tasks like SQuAD (a linear layers on top of the
        hidden-states output to compute `span start logits` and
        `span end logits`).

        https://huggingface.co/transformers/model_doc/bert.html#transformers.models.bert.modeling_bert.BertForQuestionAnswering  # noqa: E501
        """

        def __init__(self, config):

            # import pdb
            # pdb.set_trace()

            # Call the init one parent class up.
            # Otherwise, the model will be defined twice.
            BertPreTrainedModel.__init__(self, config)

            self.num_labels = config.num_labels

            # Replace `BertModel` with SparseBertModel.
            self.bert = bert_cls(config)
            self.dropout = nn.Dropout(config.hidden_dropout_prob)
            self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

            self.init_weights()

    # Rename new model class and format docstring.
    new_cls = NewBertForQuestionAnswering
    new_cls.__name__ = new_cls.__name__.replace("New", name_prefix)
    new_cls.__doc__ = new_cls.__doc__.format(name_prefix=name_prefix)

    return new_cls<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    Author Nicolas Bessi. Copyright Camptocamp SA
##############################################################################
from shapely.wkb import dumps as wkbdumps, loads as  wkbloads
from shapely.wkt import dumps as wktdumps, loads as  wktloads
from shapely.geometry import asShape
import logging

UNION_MAPPING = {'|': 'OR', '&':'AND'}

logger = logging.getLogger('GeoEngine sql debug')

# TODO Refactor geo_search and dry up the get_**_sql code

def _get_geo_func(model, domain):
    current_field = model._columns[domain[0]]
    current_operator = current_field._geo_operator
    attr = "get_%s_sql" % (domain[1],)
    if not hasattr(current_operator, attr):
        raise ValueError('Field %s does not support %s' %(current_field, domain[1]))
    func = getattr(current_operator, attr)
    return func

def geo_search(model, cursor, uid, domain=[], geo_domain=[], offset=0, limit=None, order=None, context=None):
    context = context or {}
    model.pool.get('ir.model.access').check(cursor, uid, model._name, 'read')
    query = model._where_calc(cursor, uid, domain, active_test=True, context=context)
    model._apply_ir_rules(cursor, uid, query, 'read', context=context)
    order_by = model._generate_order_by(order, query)
    from_clause, where_clause, where_clause_params = query.get_sql()
    limit_str = limit and ' limit %d' % limit or ''
    offset_str = offset and ' offset %d' % offset or ''
    where_clause_arr= []
    if where_clause and where_clause_params:
        where_clause_arr.append(where_clause)
    #geosearch where clause generation
    MODE = ''
    UNION = 'AND'
    JOIN_MODE = '%s %s'
    for domain in geo_domain:
        if isinstance(domain, basestring):
            if domain == '!':
                MODE = 'NOT'
            if domain in UNION_MAPPING.keys():
                UNION = UNION_MAPPING[domain]
        if where_clause_arr:
            where_clause_arr.append(JOIN_MODE % (MODE, UNION))
        # We start computing geo spation SQL
        if isinstance(domain, (list, tuple)):
            if isinstance(domain[2], dict):
                # We are having indirect geo_operator like (‘geom’, ‘geo_...’, {‘res.zip.poly’: [‘ids’, ‘in’, [1,2,3]] })
                ref_search = domain[2]
                rel_where_statement = []
                for key in ref_search:
                    i = key.rfind('.')
                    rel_model = key[0:i]
                    rel_col = key[i+1:]
                    rel_model = model.pool.get(rel_model)
                    from_clause += ', %s' % (rel_model._table,)
                    att_where_sql = u''
                    # we compute the attributes search on spatial rel
                    if ref_search[key]:
                        rel_query = rel_model._where_calc(cursor, uid, ref_search[key],
                                                      active_test=True, context=context)
                        rel_res = rel_query.get_sql()
                        att_where_sql = rel_res[1]
                        where_clause_params += rel_res[2]
                    # we compute the spatial search on spatial rel
                    func = _get_geo_func(model, domain)
                    spatial_where_sql = func(model._table, domain[0], domain[2],
                                             rel_col=rel_col, rel_model=rel_model)
                    if att_where_sql:
                        rel_where_statement.append(u"(%s AND %s)" % (att_where_sql, spatial_where_sql))
                    else:
                        rel_where_statement.append(u"(%s)" % (spatial_where_sql))
                where_clause_arr.append(u"AND ".join(rel_where_statement))
            else:
                current_field = model._columns[domain[0]]
                func = _get_geo_func(model, domain)
                where_sql = func(model._table, domain[0], domain[2])
                where_clause_arr.append(where_sql)
    if where_clause_arr:
        where_statement =  " WHERE %s" % (u' '.join(where_clause_arr))
    else:
        where_statement = u''
    sql= 'SELECT "%s".id FROM ' % model._table + from_clause + where_statement + order_by or '' + limit_str or '' + offset_str or ''

    logger.debug(cursor.mogrify(sql, where_clause_params))
    cursor.execute(sql, where_clause_params)
    res = cursor.fetchall()
    if res :
        return [x[0] for x in res]
    else:
        return []

class GeoOperator(object):

    def __init__(self, geo_field):
        self.geo_field = geo_field

    def get_rel_field(self, rel_col, rel_model):
        """Retriev the expression to use in PostGIS statement for a spatial
           rel search"""
        try:
            rel_model._columns[rel_col]
        except Exception, exc:
            raise Exception('Model %s has no column %s' % (rel_model._name, rel_col))
        return "%s.%s" %(rel_model._table, rel_col)

    def _get_direct_como_op_sql(self, table, col, value, rel_col=None, rel_model=None, op=''):
        "provide raw sql for geater and lesser operators"
        if isinstance(value, (int, long, float)):
            if  rel_col and rel_model:
                raise Exception('Area %s does not support int compare for relation search' % (op,))
            return " ST_Area(%s.%s) %s %s" %(table, col, op, value)
        else:
            if rel_col and rel_model:
               compare_to = self.get_rel_field(rel_col, rel_model)
            else:
                base = self.geo_field.entry_to_shape(value, same_type=False)
                compare_to = base.wkt
            return " ST_Area(%s.%s) %s ST_Area(ST_GeomFromText('%s'))" %(table, col, op, compare_to)

    def _get_postgis_comp_sql(self, table, col, value, rel_col=None, rel_model=None, op=''):
        "return raw sql for all search based on St_**(a, b) posgis operator"
        if rel_col and rel_model:
           compare_to = self.get_rel_field(rel_col, rel_model)
        else:
            base = self.geo_field.entry_to_shape(value, same_type=False)
            compare_to = "ST_GeomFromText('%s')" %(base.wkt,)
        return " %s(%s.%s, %s)" %(op, table, col, compare_to)

    ## Area comparison #############
    def get_geo_greater_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_greater operator"
        return self._get_direct_como_op_sql(table, col, value, rel_col=rel_col, rel_model=rel_model, op='>')

    def get_geo_lesser_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_lesser operator"
        return self._get_direct_como_op_sql(table, col, value, rel_col=rel_col, rel_model=rel_model, op='>')

    ## Equality comparison #############

    def get_geo_equal_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_equal operator"
        if rel_col and rel_model:
           compare_to = self.get_rel_field(rel_col, rel_model)
        else:
            base = self.geo_field.entry_to_shape(value, same_type=False)
            compare_to = "ST_GeomFromText('%s')" %(base.wkt,)
        return " %.%s = %s" %(table, col, compare_to)

    ## PostGis spatial comparison ###########

    def get_geo_intersect_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_intersec operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Intersects')

    def get_geo_touch_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_touch operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Touches')

    def get_geo_within_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_within operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Within')

<commit_msg>[ADD] geometry operator sample<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    Author Nicolas Bessi. Copyright Camptocamp SA
##############################################################################
from shapely.wkb import dumps as wkbdumps, loads as  wkbloads
from shapely.wkt import dumps as wktdumps, loads as  wktloads
from shapely.geometry import asShape
import logging

UNION_MAPPING = {'|': 'OR', '&':'AND'}

logger = logging.getLogger('GeoEngine sql debug')

# TODO Refactor geo_search and dry up the get_**_sql code

def _get_geo_func(model, domain):
    current_field = model._columns[domain[0]]
    current_operator = current_field._geo_operator
    attr = "get_%s_sql" % (domain[1],)
    if not hasattr(current_operator, attr):
        raise ValueError('Field %s does not support %s' %(current_field, domain[1]))
    func = getattr(current_operator, attr)
    return func

def geo_search(model, cursor, uid, domain=[], geo_domain=[], offset=0, limit=None, order=None, context=None):
    context = context or {}
    model.pool.get('ir.model.access').check(cursor, uid, model._name, 'read')
    query = model._where_calc(cursor, uid, domain, active_test=True, context=context)
    model._apply_ir_rules(cursor, uid, query, 'read', context=context)
    order_by = model._generate_order_by(order, query)
    from_clause, where_clause, where_clause_params = query.get_sql()
    limit_str = limit and ' limit %d' % limit or ''
    offset_str = offset and ' offset %d' % offset or ''
    where_clause_arr= []
    if where_clause and where_clause_params:
        where_clause_arr.append(where_clause)
    #geosearch where clause generation
    MODE = ''
    UNION = 'AND'
    JOIN_MODE = '%s %s'
    for domain in geo_domain:
        if isinstance(domain, basestring):
            if domain == '!':
                MODE = 'NOT'
            if domain in UNION_MAPPING.keys():
                UNION = UNION_MAPPING[domain]
        if where_clause_arr:
            where_clause_arr.append(JOIN_MODE % (MODE, UNION))
        # We start computing geo spation SQL
        if isinstance(domain, (list, tuple)):
            if isinstance(domain[2], dict):
                # We are having indirect geo_operator like (‘geom’, ‘geo_...’, {‘res.zip.poly’: [‘ids’, ‘in’, [1,2,3]] })
                ref_search = domain[2]
                rel_where_statement = []
                for key in ref_search:
                    i = key.rfind('.')
                    rel_model = key[0:i]
                    rel_col = key[i+1:]
                    rel_model = model.pool.get(rel_model)
                    from_clause += ', %s' % (rel_model._table,)
                    att_where_sql = u''
                    # we compute the attributes search on spatial rel
                    if ref_search[key]:
                        rel_query = rel_model._where_calc(cursor, uid, ref_search[key],
                                                      active_test=True, context=context)
                        rel_res = rel_query.get_sql()
                        att_where_sql = rel_res[1]
                        where_clause_params += rel_res[2]
                    # we compute the spatial search on spatial rel
                    func = _get_geo_func(model, domain)
                    spatial_where_sql = func(model._table, domain[0], domain[2],
                                             rel_col=rel_col, rel_model=rel_model)
                    if att_where_sql:
                        rel_where_statement.append(u"(%s AND %s)" % (att_where_sql, spatial_where_sql))
                    else:
                        rel_where_statement.append(u"(%s)" % (spatial_where_sql))
                where_clause_arr.append(u"AND ".join(rel_where_statement))
            else:
                current_field = model._columns[domain[0]]
                func = _get_geo_func(model, domain)
                where_sql = func(model._table, domain[0], domain[2])
                where_clause_arr.append(where_sql)
    if where_clause_arr:
        where_statement =  " WHERE %s" % (u' '.join(where_clause_arr))
    else:
        where_statement = u''
    sql= 'SELECT "%s".id FROM ' % model._table + from_clause + where_statement + order_by or '' + limit_str or '' + offset_str or ''
    print cursor.mogrify(sql, where_clause_params)
    logger.debug(cursor.mogrify(sql, where_clause_params))
    cursor.execute(sql, where_clause_params)
    res = cursor.fetchall()
    if res :
        return [x[0] for x in res]
    else:
        return []

class GeoOperator(object):

    def __init__(self, geo_field):
        self.geo_field = geo_field

    def get_rel_field(self, rel_col, rel_model):
        """Retriev the expression to use in PostGIS statement for a spatial
           rel search"""
        try:
            rel_model._columns[rel_col]
        except Exception, exc:
            raise Exception('Model %s has no column %s' % (rel_model._name, rel_col))
        return "%s.%s" %(rel_model._table, rel_col)

    def _get_direct_como_op_sql(self, table, col, value, rel_col=None, rel_model=None, op=''):
        "provide raw sql for geater and lesser operators"
        if isinstance(value, (int, long, float)):
            if  rel_col and rel_model:
                raise Exception('Area %s does not support int compare for relation search' % (op,))
            return " ST_Area(%s.%s) %s %s" %(table, col, op, value)
        else:
            if rel_col and rel_model:
               compare_to = self.get_rel_field(rel_col, rel_model)
            else:
                base = self.geo_field.entry_to_shape(value, same_type=False)
                compare_to = base.wkt
            return " ST_Area(%s.%s) %s ST_Area(ST_GeomFromText('%s'))" %(table, col, op, compare_to)

    def _get_postgis_comp_sql(self, table, col, value, rel_col=None, rel_model=None, op=''):
        "return raw sql for all search based on St_**(a, b) posgis operator"
        if rel_col and rel_model:
           compare_to = self.get_rel_field(rel_col, rel_model)
        else:
            base = self.geo_field.entry_to_shape(value, same_type=False)
            compare_to = "ST_GeomFromText('%s')" %(base.wkt,)
        return " %s(%s.%s, %s)" %(op, table, col, compare_to)

    ## Area comparison #############
    def get_geo_greater_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_greater operator"
        return self._get_direct_como_op_sql(table, col, value, rel_col=rel_col, rel_model=rel_model, op='>')

    def get_geo_lesser_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_lesser operator"
        return self._get_direct_como_op_sql(table, col, value, rel_col=rel_col, rel_model=rel_model, op='>')

    ## Equality comparison #############

    def get_geo_equal_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_equal operator"
        if rel_col and rel_model:
           compare_to = self.get_rel_field(rel_col, rel_model)
        else:
            base = self.geo_field.entry_to_shape(value, same_type=False)
            compare_to = "ST_GeomFromText('%s')" %(base.wkt,)
        return " %.%s = %s" %(table, col, compare_to)

    ## PostGis spatial comparison ###########

    def get_geo_intersect_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_intersec operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Intersects')

    def get_geo_touch_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_touch operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Touches')

    def get_geo_within_sql(self, table, col, value, rel_col=None, rel_model=None):
        "return raw sql for geo_within operator"
        return self._get_postgis_comp_sql(table, col, value, rel_col= rel_col, rel_model= rel_model, op='ST_Within')

<|endoftext|>


<commit_before>from datetime import datetime
import glob
from optparse import make_option
import os
import sys
from time import sleep

from django.core.management.base import BaseCommand, CommandError
from django.conf import settings

from eulcore.xmlmap.core import load_xmlobject_from_file
from eulcore.django.existdb import ExistDB, ExistDBException

from findingaids.fa.models import FindingAid
from findingaids.fa_admin.utils import check_ead
from findingaids.fa_admin.tasks import reload_cached_pdf

class Command(BaseCommand):
    """Load all or specified EAD xml files in the configured source directory
to the configured eXist collection.  For each document successfully loaded to
eXist, this script will trigger a celery task to reload the PDF in the cache; the
script will not exit until all tasks have completed.

If filenames are specified as arguments, only those files will be loaded.
Files should be specified by basename only (they will be loaded from the configured
EAD source directory). Otherwise, all .xml files in the configured EAD source
directory will be loaded."""
    help = __doc__

    args = '[<filename filename ... >]'

    option_list = BaseCommand.option_list + (
        make_option('--skip-pdf-reload', '-s',
            action='store_true',
            dest='skip_pdf_reload',
            help='Skip reloading PDFs in the cache.'),
        )

    def handle(self, *args, **options):
        verbosity = int(options['verbosity'])    # 1 = normal, 0 = minimal, 2 = all
        v_normal = 1
        v_all = 2

        # check for required settings
        if not hasattr(settings, 'EXISTDB_ROOT_COLLECTION') or not settings.EXISTDB_ROOT_COLLECTION:
            raise CommandError("EXISTDB_ROOT_COLLECTION setting is missing")
            return
        if not hasattr(settings, 'FINDINGAID_EAD_SOURCE') or not settings.FINDINGAID_EAD_SOURCE:
            raise CommandError("FINDINGAID_EAD_SOURCE setting is missing")
            return

        if verbosity == v_all:
            print "Loading documents from configured EAD source directory: %s" \
                    % settings.FINDINGAID_EAD_SOURCE
            print 'Documents will be loaded to configured eXist collection: %s' \
                    % settings.EXISTDB_ROOT_COLLECTION
            if options['skip_pdf_reload']:
                print "** Skipping PDFs cache reload"

        db = ExistDB()

        loaded = 0
        errored = 0
        pdf_tasks = {}

        start_time = datetime.now()
        try:
            if len(args):
                files = [os.path.join(settings.FINDINGAID_EAD_SOURCE, name) for name in args]
            else:
                files = glob.iglob(os.path.join(settings.FINDINGAID_EAD_SOURCE, '*.xml'))
                self.db = ExistDB()

            for file in files:
                try:
                    # full path location where file will be loaded in exist db collection
                    dbpath = settings.EXISTDB_ROOT_COLLECTION + "/" + os.path.basename(file)
                    errors = check_ead(file, dbpath)
                    if errors:
                        # report errors, don't load
                        errored += 1
                        print "Error: %s does not pass publication checks; not loading to eXist." % file
                        if verbosity >= v_normal:
                            print "  Errors found:"
                            for err in errors:
                                print "    %s" % err
                    else:
                        with open(file, 'r') as eadfile:
                            success = db.load(eadfile, dbpath, overwrite=True)

                        if success:
                            loaded += 1
                            if verbosity >= v_normal:
                                print "Loaded %s" % file
                            # load the file as a FindingAid object to get the eadid for PDF reload
                            ead = load_xmlobject_from_file(file, FindingAid)

                            # trigger PDF regeneration in the cache and store task result
                            # - unless user has requested PDF reload be skipped
                            if not options['skip_pdf_reload']:
                                pdf_tasks[ead.eadid] = reload_cached_pdf.delay(ead.eadid.value)
                                # NOTE: unlike the web admin publish, this does not
                                # generate TaskResult db records; task outcomes will be
                                # checked & reported before the script finishes
                        else:
                            errored += 1
                            print "Error: failed to load %s to eXist" % file
                except ExistDBException, e:
                    print "Error: failed to load %s to eXist" % file
                    print e.message()
                    errored += 1

            # output a summary of what was done
            print "%d document%s loaded" % (loaded, 's' if loaded != 1 else '')
            print "%d document%s with errors" % (errored, 's' if errored != 1 else '')

            if not options['skip_pdf_reload']:
                # check on the status of PDF cache reload tasks and wait until they all finish
                success, failed, pending = check_tasks(pdf_tasks)
                msg = ''
                while pending:
                    if verbosity >= v_normal:
                        # back up the length of the last message so it can be overwritten
                        for i in range(len(msg)):
                            sys.stdout.write('\r')
                        msg = "Waiting for %d PDF cache reload task%s to complete..." % \
                                (len(pending), 's' if len(pending) != 1 else '')
                        sys.stdout.write(msg)
                        sys.stdout.flush()
                    sleep(5)
                    # check the remaining incomplete tasks, updating success/failure lists
                    s, f, pending = check_tasks(pending)
                    success.update(s)
                    failed.update(f)

                if verbosity >= v_normal:
                    # remove any last pending output message
                    for i in range(len(msg)):
                        sys.stdout.write('\r')
                    print ''    # print a newline after any pending output

                print "Successfully reloaded PDFs for %d document%s" % \
                        (len(success), 's' if len(success) != 1 else '')
                if verbosity >= v_all:
                    print ', '.join(success.keys())
                print "Failed to reloaded PDFs for %d document%s" % \
                        (len(failed), 's' if len(failed) != 1 else '')
                print ', '.join(failed.keys())

            end_time = datetime.now()
            if verbosity >= v_normal:
                print "Ran for %s" % str(end_time - start_time)

        except Exception, err:
            raise CommandError(err)

def check_tasks(tasks):
    """Check the status of celery tasks for successful completion.  Expects
    a dictionary with values that are instances of :class:`celery.result.AsyncResult`.

    :returns: three dictionaries with the same keys and AsyncResult values as the
        input parameter, in this order: succeeded (task finished successfully),
        failed (task finished unsuccessfully), and pending (task not yet finished)
    """
    succeeded = {}
    failed = {}
    pending = {}
    for id, task in tasks.iteritems():
        if task.ready():        # task has finished running
            if task.successful():  # SUCCESS
                succeeded[id] = task
            else:   # FAILURE
                failed[id] = task
        else:   # not yet finished
            pending[id] = task
    return succeeded, failed, pending<commit_msg>use new eadid format for tracking pdf reload tasks<commit_after>from datetime import datetime
import glob
from optparse import make_option
import os
import sys
from time import sleep

from django.core.management.base import BaseCommand, CommandError
from django.conf import settings

from eulcore.xmlmap.core import load_xmlobject_from_file
from eulcore.django.existdb import ExistDB, ExistDBException

from findingaids.fa.models import FindingAid
from findingaids.fa_admin.utils import check_ead
from findingaids.fa_admin.tasks import reload_cached_pdf

class Command(BaseCommand):
    """Load all or specified EAD xml files in the configured source directory
to the configured eXist collection.  For each document successfully loaded to
eXist, this script will trigger a celery task to reload the PDF in the cache; the
script will not exit until all tasks have completed.

If filenames are specified as arguments, only those files will be loaded.
Files should be specified by basename only (they will be loaded from the configured
EAD source directory). Otherwise, all .xml files in the configured EAD source
directory will be loaded."""
    help = __doc__

    args = '[<filename filename ... >]'

    option_list = BaseCommand.option_list + (
        make_option('--skip-pdf-reload', '-s',
            action='store_true',
            dest='skip_pdf_reload',
            help='Skip reloading PDFs in the cache.'),
        )

    def handle(self, *args, **options):
        verbosity = int(options['verbosity'])    # 1 = normal, 0 = minimal, 2 = all
        v_normal = 1
        v_all = 2

        # check for required settings
        if not hasattr(settings, 'EXISTDB_ROOT_COLLECTION') or not settings.EXISTDB_ROOT_COLLECTION:
            raise CommandError("EXISTDB_ROOT_COLLECTION setting is missing")
            return
        if not hasattr(settings, 'FINDINGAID_EAD_SOURCE') or not settings.FINDINGAID_EAD_SOURCE:
            raise CommandError("FINDINGAID_EAD_SOURCE setting is missing")
            return

        if verbosity == v_all:
            print "Loading documents from configured EAD source directory: %s" \
                    % settings.FINDINGAID_EAD_SOURCE
            print 'Documents will be loaded to configured eXist collection: %s' \
                    % settings.EXISTDB_ROOT_COLLECTION
            if options['skip_pdf_reload']:
                print "** Skipping PDFs cache reload"

        db = ExistDB()

        loaded = 0
        errored = 0
        pdf_tasks = {}

        start_time = datetime.now()
        try:
            if len(args):
                files = [os.path.join(settings.FINDINGAID_EAD_SOURCE, name) for name in args]
            else:
                files = glob.iglob(os.path.join(settings.FINDINGAID_EAD_SOURCE, '*.xml'))
                self.db = ExistDB()

            for file in files:
                try:
                    # full path location where file will be loaded in exist db collection
                    dbpath = settings.EXISTDB_ROOT_COLLECTION + "/" + os.path.basename(file)
                    errors = check_ead(file, dbpath)
                    if errors:
                        # report errors, don't load
                        errored += 1
                        print "Error: %s does not pass publication checks; not loading to eXist." % file
                        if verbosity >= v_normal:
                            print "  Errors found:"
                            for err in errors:
                                print "    %s" % err
                    else:
                        with open(file, 'r') as eadfile:
                            success = db.load(eadfile, dbpath, overwrite=True)

                        if success:
                            loaded += 1
                            if verbosity >= v_normal:
                                print "Loaded %s" % file
                            # load the file as a FindingAid object to get the eadid for PDF reload
                            ead = load_xmlobject_from_file(file, FindingAid)

                            # trigger PDF regeneration in the cache and store task result
                            # - unless user has requested PDF reload be skipped
                            if not options['skip_pdf_reload']:
                                pdf_tasks[ead.eadid.value] = reload_cached_pdf.delay(ead.eadid.value)
                                # NOTE: unlike the web admin publish, this does not
                                # generate TaskResult db records; task outcomes will be
                                # checked & reported before the script finishes
                        else:
                            errored += 1
                            print "Error: failed to load %s to eXist" % file
                except ExistDBException, e:
                    print "Error: failed to load %s to eXist" % file
                    print e.message()
                    errored += 1

            # output a summary of what was done
            print "%d document%s loaded" % (loaded, 's' if loaded != 1 else '')
            print "%d document%s with errors" % (errored, 's' if errored != 1 else '')

            if not options['skip_pdf_reload']:
                # check on the status of PDF cache reload tasks and wait until they all finish
                success, failed, pending = check_tasks(pdf_tasks)
                msg = ''
                while pending:
                    if verbosity >= v_normal:
                        # back up the length of the last message so it can be overwritten
                        for i in range(len(msg)):
                            sys.stdout.write('\r')
                        msg = "Waiting for %d PDF cache reload task%s to complete..." % \
                                (len(pending), 's' if len(pending) != 1 else '')
                        sys.stdout.write(msg)
                        sys.stdout.flush()
                    sleep(5)
                    # check the remaining incomplete tasks, updating success/failure lists
                    s, f, pending = check_tasks(pending)
                    success.update(s)
                    failed.update(f)

                if verbosity >= v_normal:
                    # remove any last pending output message
                    for i in range(len(msg)):
                        sys.stdout.write('\r')
                    print ''    # print a newline after any pending output

                print "Successfully reloaded PDFs for %d document%s" % \
                        (len(success), 's' if len(success) != 1 else '')
                if verbosity >= v_all:
                    print ', '.join(success.keys())
                print "Failed to reloaded PDFs for %d document%s" % \
                        (len(failed), 's' if len(failed) != 1 else '')
                print ', '.join(failed.keys())

            end_time = datetime.now()
            if verbosity >= v_normal:
                print "Ran for %s" % str(end_time - start_time)

        except Exception, err:
            raise CommandError(err)

def check_tasks(tasks):
    """Check the status of celery tasks for successful completion.  Expects
    a dictionary with values that are instances of :class:`celery.result.AsyncResult`.

    :returns: three dictionaries with the same keys and AsyncResult values as the
        input parameter, in this order: succeeded (task finished successfully),
        failed (task finished unsuccessfully), and pending (task not yet finished)
    """
    succeeded = {}
    failed = {}
    pending = {}
    for id, task in tasks.iteritems():
        if task.ready():        # task has finished running
            if task.successful():  # SUCCESS
                succeeded[id] = task
            else:   # FAILURE
                failed[id] = task
        else:   # not yet finished
            pending[id] = task
    return succeeded, failed, pending<|endoftext|>