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

# This is the refactored version for the benchmark script. Once this
# is completely under test coverage, the original benchmark.py gets
# turned into this, then this can be deleted.

# TODO: Why are records over & underclumped? ie. Charles Linder 7475467,
# Nirmala Ramanujam 7570988, Edwin L. Thomas 7799416

# TODO: Investigate "WEIRDNESS" output : Jaro-Winkler exception
# thrown on comparison between  and ROBERT BRUCE
# Underclumping: unique records vs. total records

import sqlite3, sys, csv, datetime;
import argparse
sys.path.append( './lib/' )
sys.path.append("lib")
from fwork import *
from bmconfig import *

# This is what is known as "testable code"
def get_filename_suffix(filename):
    return filename.split(".")[-1].lower()

# More testable code...
def is_csv_file(filename):
    return get_filename_suffix(filename) == "csv"

def print_diagnostics(data, table, header, tList):
   print "Printing diagnostics..."
   print "data: ", data
   print "table: ", table
   print "header: ", header
   print "tList: ", tList

def create_match_tables(c, fBnme, uqB, exCom, exAnd):
    # TODO: Split this query into two functions, test each
    c.executescript("""
        /* EXPAND UNIQUE BASE AND INDICATE ACTIVE MATCHES */
        CREATE TABLE dataM3 AS
        SELECT  uqS, a.*
          FROM (SELECT  uqS AS uqSUB, a.*
              FROM (SELECT  uqB, b.*
                  FROM  (SELECT DISTINCT(uqB) FROM dataM2 WHERE uqB!="") AS a
                INNER JOIN  %s AS b
                    ON  a.uqB=b.%s) AS a
             LEFT JOIN (SELECT %s, uqB, uqS FROM dataM2) AS b
                ON  a.uqB=b.uqB AND %s) AS a
        INNER JOIN (SELECT DISTINCT uqB, uqS FROM dataM2) AS b
            ON  a.%s=b.uqB;

        /* INDICATE INVENTORS WHO DO NOT MATCH */
        CREATE TABLE dataM4 AS
        SELECT  errD(a.ErrUQ, uqB) AS ErrUQ, b.*
          FROM (SELECT uqS, freqUQ(uqB) as ErrUQ FROM dataM3 GROUP BY uqS) AS a
        INNER JOIN  dataM3 AS b
            ON  a.uqS=b.uqS AND b.AppYear <= '2010' /*AND a.uqS not in (83, 85, 93)*/
          ORDER BY  uqS, %s;

        """ % (fBnme, uqB, exCom, exAnd, uqB, exCom))

def handle_fuzzy_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd):
	# TODO: Remove leading CREATE INDEX as its already been created in the
	# calling function
	# TODO: Split this into 3 functions, no reason to do all this work in
	# one monster query.
        c.executescript("""
        CREATE INDEX IF NOT EXISTS dS_E ON dataS (%s);

        /* RETAIN ONLY JARO>0.9 FUZZY AND EXACT MATCHES */
        CREATE TABLE dataM AS
            SELECT  a.*, %s AS uqB, %s AS uqS, %s AS jaro
              FROM  %s AS a
        INNER JOIN  dataS AS b
            ON  %s
             WHERE  jaro>0.90;

        /* DETERMINE MAXIMUM JARO FOR EACH UQ AND EXACT COMBO */
        CREATE TABLE dataT AS
            SELECT  uqS, %s, MAX(jaro) AS jaro, count(*) as cnt
              FROM  dataM
          GROUP BY  uqS, %s;

        /* RETAIN ONLY MAXIMUM JARO */
        CREATE TABLE dataM2 AS
            SELECT  a.*
              FROM  dataM AS a
        INNER JOIN  dataT AS b
            ON  a.uqS=b.uqS AND a.jaro=b.jaro AND %s;
        """ % (exCom, uqB, uqS,
               "*".join(["jarow(a.%s, b.%s)" % (x,x) for x in fuzzy]),
               fBnme, exAnd, exCom, exCom, exAnd))

def handle_nonfuzzy_dataS(uqB, uqS, fBnme, exAnd):
    c.executescript("""
	CREATE TABLE dataM2 AS
	    SELECT  *, %s AS uqB, %s AS uqS
	      FROM  %s AS a
	INNER JOIN  dataS AS b
		ON  %s;
	""" % (uqB, uqS, fBnme, exAnd))

def export_csv_results(c, output):
	writer = csv.writer(open(output, "wb"), lineterminator="\n")
	writer.writerows([[x[1] for x in c.execute("PRAGMA TABLE_INFO(dataM4)")]])
	writer.writerows(c.execute("SELECT * FROM dataM4").fetchall())

# TODO: Create functions for handling true and false positives and negatives,
# compute all the relevant statistics using those measures instead of the
# mess of inline computation following the heredoc.
# TODO: Switch to printing json results
def print_results(c, output, t):
	#print "Printing results ..." + str(datetime.datetime.now())
	rep = [list(x) for x in c.execute("SELECT ErrUQ, uqSUB FROM dataM4")]
	#print "Rep: ", rep
	orig = len([x for x in rep if x[1]!=None])
	errm = sum([int(x[0]) for x in rep if x[0]!=None])
	#print errm
	u = 1.0*errm/orig
	o = 1-(float(orig)/len(rep))
	recall = 1.0 - u
	# overclumping is lumping
	# underclumping is splitting
	print """

	RESULTS ==================

	     Original: {original}
	  New Records: {new}
		Total: {total}

	    Overclump: {overclump} ({o:.2%})
	   Underclump: {underclump} ({u:.2%})
	    Precision: {precision:.2%}
	       Recall: {recall:.2%}
	  File Detail: {filename}
		 Time: {time}
	""".format(original = orig, new = len(rep)-orig, total = len(rep), overclump = len(rep)-orig, o = o,
		   underclump = errm, u = u, recall = recall, precision = recall/(recall+o), filename = output,
		   time = datetime.datetime.now()-t)

def attach_database(c, fileB, tblB, exCom, exAnd):
	# TODO: Replace with call to is_csv_file(fileB),
	# after it's unit tested.
	if is_csv_file(fileB): #fileB.split(".")[-1].lower()=="csv":
	# TODO: Try to move some of this to a function
	    dataB = uniVert([x for x in csv.reader(open("%s" % fileB, "rb"))])

	    print_diagnostics(dataB, "dataB", True, ["Patent VARCHAR"])
	    #quickSQL(c, data=dataB,  table="dataB", header=True, typeList=["Patent VARCHAR"])
	    quickSQL2(c, data=dataB,  table="dataB", header=True, typeList=["Patent VARCHAR"])

	    c.execute("CREATE INDEX IF NOT EXISTS dB_E ON dataB (%s)" % exCom)
	    c.execute("CREATE INDEX IF NOT EXISTS dB_U ON dataB (%s)" % uqB)
	    fBnme = "dataB"
	# else assume fileB is an SQLite3 database file...
	else:
	    c.execute("ATTACH DATABASE '%s' AS db" % fileB)
	    # fBnme is, apparently a table name. Or maybe a tbl nme...
	    if tblB=="":
		fBnme = "db.%s" % fileB.split(".")[-2].split("/")[-1]
	    else:
		fBnme = "db.%s" % tblB

        return fBnme

def handle_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd):
	c.execute("CREATE INDEX IF NOT EXISTS dS_E ON dataS (%s);" % (exCom))
	if fuzzy:
	    handle_fuzzy_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd)
	else:
	    handle_nonfuzzy_dataS(uqB, uqS, fBnme, exAnd)

def print_match_qualifiers(exact, fuzzy, uqS):
	print "Exact: ", exact
	print "Fuzzy: ", fuzzy
	print "uqS: ", uqS

def bmVerify(results, filepath="", outdir = ""):
        """
        Analysis function on disambiguation results, assuming that all benchmark data
        are in the large results dataset.

        Creates analysis detail csv file and prints summary information on
        over- and underclumping statistics.

        Running from the command line (make sure to set correct file paths in file)
        python bmVerify_v3.py "input filepath" "output directory" databases
        example:
        python bmVerify_v3.py /home/ysun/disambig/newcode/all/ /home/ayu/results_v2/ invpatC_NBNA.good.Jan2011 invpatC_NBYA.good.Jan2011 invpatC_YBNA.good.Jan2011

        Running interactively:
        import bmVerify_v3
        bmVerify(['final_r7', 'final_r8'], filepath="/home/ysun/disambig/newcode/all/", outdir = "/home/ayu/results_v2/")

        """

	# TODO: Move these out of this function completely, and into the top
	# part of this file.
	uqB = "Unique_Inventor_ID2"
	tblB = "invpat"
	#fileS = "/home/ron/disambig/BM/DefTruth5.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/DefTruth5.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/berkeley.csv"
	fileS = "/home/doolin/src/patentbenchmarks/pister.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/paulrgray.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/allbritton.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/siddhu.csv"
	#fileS = "/var/share/patentdata/disambiguation/experiments/earth/berkeley/benchmark.csv"

        for result in results:

                fileB = filepath + "{result}.sqlite3".format(result=result)
                output = outdir + "{result}_DT5.csv".format(result=result)

                t = datetime.datetime.now()

                #print "Start time: " + str(datetime.datetime.now())

		# TODO: Move freqUQ out of this function if possible.
                class freqUQ:
                    def __init__(self):
                        self.list=[]
                    def step(self, value):
                        self.list.append(value)
                    def finalize(self):
                        return sorted([(self.list.count(x), x) for x in set(self.list)], reverse=True)[0][1]

                #MAKE THIS SO IT CAN ATTACH SQLITE3 FOR BENCHMARK
		# UniqueID,Patent,Lastname,Firstname
		# ,VARCHAR,,
		# ,%08d,,
		# UNIQUE,EXACT,FUZZY,FUZZY
		# 1,5773227,ALLBRITTON,NANCY L
		# TODO: Refactor into read_benchmark_csvfile()
		# This will help for unit test and provide self-documentation
                dataS = uniVert([x for x in csv.reader(open(fileS, "rb"))])
                #print dataS

		# TODO: Refactor into create_column_types()
                #1 = Variables, 2 = Type, 3 = Format (If necessary), 4 = Matching Type
                tList = ["%s %s" % (dataS[0][i], x) for i,x in enumerate(dataS[1]) if  x != ""]

		# Slice out rows 1, 2 & 3 from dataS. This is the data which gets
		# put in the database created below, and used for matching.
		# TODO: Refactor into slice_header_from()
                dataS2 = [dataS[0]]
                dataS2.extend(dataS[4:])

                # Format if its necessary --> Basically for Patents..
                for i,x in enumerate(dataS[2]):
                    if x!="":
                        for j in xrange(1,len(dataS2)):
                            if dataS2[j][i].isdigit():
                                dataS2[j][i] = x % int(dataS2[j][i])

                conn = sqlite3.connect(":memory:")
                #conn = sqlite3.connect("bmark2.sqlite3")
                conn.create_function("jarow", 2, jarow)
                conn.create_function("errD", 2, lambda x,y: (x!=y) and 1 or None)
                conn.create_aggregate("freqUQ", 1, freqUQ)
                c = conn.cursor()

                # FIGURE OUT WHICH ONES HAVE EXACT/FUZZY
                exact = [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="E"]
                fuzzy = [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="F"]
                uqS =   [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="U"][0]
                #print_match_qualifiers(exact, fuzzy, uqS)

                # CREATE INDEX, MERGE DATA BASED ON EXACTS
                #print "Creating indices... " + str(datetime.datetime.now())

                exAnd = " AND ".join(["a.%s=b.%s" % (x, x) for x in exact])
		#print "exAnd: ", exAnd
                exCom = ", ".join(exact)
                #print "exCom: ", exCom

                fBnme = attach_database(c, fileB, tblB, exCom, exAnd)

                #print_diagnostics(dataS2, "dataS", True, tList)
                #quickSQL(c, data=dataS2, table="dataS", header=True, typeList=tList)
                quickSQL2(c, data=dataS2, table="dataS", header=True, typeList=tList)

		handle_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd)

                create_match_tables(c, fBnme, uqB, exCom, exAnd)

                #print "Indices Done ... " + str(datetime.datetime.now())

		export_csv_results(c, output)
		print_results(c, output, t)
                c.close()
                conn.close()

if __name__ == "__main__":
    import sys
    if(sys.argv[1] == 'help' or sys.argv[1] == '?'):
        print bmVerify.__doc__

    else:
            bmVerify(sys.argv[3:], sys.argv[1], sys.argv[2])
<commit_msg>some cleanup on sql formatting<commit_after>#!/usr/bin/env python

# This is the refactored version for the benchmark script. Once this
# is completely under test coverage, the original benchmark.py gets
# turned into this, then this can be deleted.

# TODO: Why are records over & underclumped? ie. Charles Linder 7475467,
# Nirmala Ramanujam 7570988, Edwin L. Thomas 7799416

# TODO: Investigate "WEIRDNESS" output : Jaro-Winkler exception
# thrown on comparison between  and ROBERT BRUCE
# Underclumping: unique records vs. total records

import sqlite3, sys, csv, datetime;
import argparse
sys.path.append( './lib/' )
sys.path.append("lib")
from fwork import *
from bmconfig import *

# This is what is known as "testable code"
def get_filename_suffix(filename):
    return filename.split(".")[-1].lower()

# More testable code...
def is_csv_file(filename):
    return get_filename_suffix(filename) == "csv"

def print_diagnostics(data, table, header, tList):
   print "Printing diagnostics..."
   print "data: ", data
   print "table: ", table
   print "header: ", header
   print "tList: ", tList

def create_match_tables(c, fBnme, uqB, exCom, exAnd):
    # TODO: Split this query into two functions, test each
    c.executescript("""
        /* EXPAND UNIQUE BASE AND INDICATE ACTIVE MATCHES */
        CREATE TABLE dataM3 AS
        SELECT  uqS, a.*
          FROM (SELECT  uqS AS uqSUB, a.*
              FROM (SELECT  uqB, b.*
                  FROM  (SELECT DISTINCT(uqB) FROM dataM2 WHERE uqB!="") AS a
                INNER JOIN  %s AS b
                    ON  a.uqB=b.%s) AS a
             LEFT JOIN (SELECT %s, uqB, uqS FROM dataM2) AS b
                ON  a.uqB=b.uqB AND %s) AS a
        INNER JOIN (SELECT DISTINCT uqB, uqS FROM dataM2) AS b
            ON  a.%s=b.uqB;

        /* INDICATE INVENTORS WHO DO NOT MATCH */
        CREATE TABLE dataM4 AS
        SELECT  errD(a.ErrUQ, uqB) AS ErrUQ, b.*
          FROM (SELECT uqS, freqUQ(uqB) as ErrUQ FROM dataM3 GROUP BY uqS) AS a
        INNER JOIN  dataM3 AS b
            ON  a.uqS=b.uqS AND b.AppYear <= '2010' /*AND a.uqS not in (83, 85, 93)*/
          ORDER BY  uqS, %s;

        """ % (fBnme, uqB, exCom, exAnd, uqB, exCom))

def handle_fuzzy_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd):
	# TODO: Remove leading CREATE INDEX as its already been created in the
	# calling function
	# TODO: Split this into 3 functions, no reason to do all this work in
	# one monster query.
        c.executescript("""
               CREATE INDEX
	      IF NOT EXISTS dS_E ON dataS (%s);

              /* RETAIN ONLY JARO>0.9 FUZZY AND EXACT MATCHES */
               CREATE TABLE  dataM AS
                     SELECT  a.*, %s AS uqB, %s AS uqS, %s AS jaro
                       FROM  %s AS a
                 INNER JOIN  dataS AS b
                         ON  %s
                      WHERE  jaro>0.90;

              /* DETERMINE MAXIMUM JARO FOR EACH UQ AND EXACT COMBO */
               CREATE TABLE  dataT AS
                     SELECT  uqS, %s, MAX(jaro)
		         AS  jaro, count(*) as cnt
                       FROM  dataM
                   GROUP BY  uqS, %s;

              /* RETAIN ONLY  MAXIMUM JARO */
               CREATE TABLE  dataM2 AS
                     SELECT  a.*
                       FROM  dataM
		         AS  a
                 INNER JOIN  dataT AS b
                         ON  a.uqS=b.uqS
                        AND  a.jaro=b.jaro
			AND  %s;
                        """  % (exCom, uqB, uqS,
                               "*".join(["jarow(a.%s, b.%s)" % (x,x) for x in fuzzy]),
                                fBnme, exAnd, exCom, exCom, exAnd))

def handle_nonfuzzy_dataS(uqB, uqS, fBnme, exAnd):
    c.executescript("""
	CREATE TABLE dataM2 AS
	    SELECT  *, %s AS uqB, %s AS uqS
	      FROM  %s AS a
	INNER JOIN  dataS AS b
		ON  %s;
	""" % (uqB, uqS, fBnme, exAnd))

def export_csv_results(c, output):
	writer = csv.writer(open(output, "wb"), lineterminator="\n")
	writer.writerows([[x[1] for x in c.execute("PRAGMA TABLE_INFO(dataM4)")]])
	writer.writerows(c.execute("SELECT * FROM dataM4").fetchall())

# TODO: Create functions for handling true and false positives and negatives,
# compute all the relevant statistics using those measures instead of the
# mess of inline computation following the heredoc.
# TODO: Switch to printing json results
def print_results(c, output, t):
	#print "Printing results ..." + str(datetime.datetime.now())
	rep = [list(x) for x in c.execute("SELECT ErrUQ, uqSUB FROM dataM4")]
	#print "Rep: ", rep
	orig = len([x for x in rep if x[1]!=None])
	errm = sum([int(x[0]) for x in rep if x[0]!=None])
	#print errm
	u = 1.0*errm/orig
	o = 1-(float(orig)/len(rep))
	recall = 1.0 - u
	# overclumping is lumping
	# underclumping is splitting
	print """

	RESULTS ==================

	     Original: {original}
	  New Records: {new}
		Total: {total}

	    Overclump: {overclump} ({o:.2%})
	   Underclump: {underclump} ({u:.2%})
	    Precision: {precision:.2%}
	       Recall: {recall:.2%}
	  File Detail: {filename}
		 Time: {time}
	""".format(original = orig, new = len(rep)-orig, total = len(rep), overclump = len(rep)-orig, o = o,
		   underclump = errm, u = u, recall = recall, precision = recall/(recall+o), filename = output,
		   time = datetime.datetime.now()-t)

def attach_database(c, fileB, tblB, exCom, exAnd):
	# TODO: Replace with call to is_csv_file(fileB),
	# after it's unit tested.
	if is_csv_file(fileB): #fileB.split(".")[-1].lower()=="csv":
	# TODO: Try to move some of this to a function
	    dataB = uniVert([x for x in csv.reader(open("%s" % fileB, "rb"))])

	    print_diagnostics(dataB, "dataB", True, ["Patent VARCHAR"])
	    #quickSQL(c, data=dataB,  table="dataB", header=True, typeList=["Patent VARCHAR"])
	    quickSQL2(c, data=dataB,  table="dataB", header=True, typeList=["Patent VARCHAR"])

	    c.execute("CREATE INDEX IF NOT EXISTS dB_E ON dataB (%s)" % exCom)
	    c.execute("CREATE INDEX IF NOT EXISTS dB_U ON dataB (%s)" % uqB)
	    fBnme = "dataB"
	# else assume fileB is an SQLite3 database file...
	else:
	    c.execute("ATTACH DATABASE '%s' AS db" % fileB)
	    # fBnme is, apparently a table name. Or maybe a tbl nme...
	    if tblB=="":
		fBnme = "db.%s" % fileB.split(".")[-2].split("/")[-1]
	    else:
		fBnme = "db.%s" % tblB

        return fBnme

def handle_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd):
	c.execute("CREATE INDEX IF NOT EXISTS dS_E ON dataS (%s);" % (exCom))
	if fuzzy:
	    handle_fuzzy_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd)
	else:
	    handle_nonfuzzy_dataS(uqB, uqS, fBnme, exAnd)

def print_match_qualifiers(exact, fuzzy, uqS):
	print "Exact: ", exact
	print "Fuzzy: ", fuzzy
	print "uqS: ", uqS

def bmVerify(results, filepath="", outdir = ""):
        """
        Analysis function on disambiguation results, assuming that all benchmark data
        are in the large results dataset.

        Creates analysis detail csv file and prints summary information on
        over- and underclumping statistics.

        Running from the command line (make sure to set correct file paths in file)
        python bmVerify_v3.py "input filepath" "output directory" databases
        example:
        python bmVerify_v3.py /home/ysun/disambig/newcode/all/ /home/ayu/results_v2/ invpatC_NBNA.good.Jan2011 invpatC_NBYA.good.Jan2011 invpatC_YBNA.good.Jan2011

        Running interactively:
        import bmVerify_v3
        bmVerify(['final_r7', 'final_r8'], filepath="/home/ysun/disambig/newcode/all/", outdir = "/home/ayu/results_v2/")

        """

	# TODO: Move these out of this function completely, and into the top
	# part of this file.
	uqB = "Unique_Inventor_ID2"
	tblB = "invpat"
	#fileS = "/home/ron/disambig/BM/DefTruth5.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/DefTruth5.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/berkeley.csv"
	fileS = "/home/doolin/src/patentbenchmarks/pister.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/paulrgray.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/allbritton.csv"
	#fileS = "/home/doolin/src/patentbenchmarks/siddhu.csv"
	#fileS = "/var/share/patentdata/disambiguation/experiments/earth/berkeley/benchmark.csv"

        for result in results:

                fileB = filepath + "{result}.sqlite3".format(result=result)
                output = outdir + "{result}_DT5.csv".format(result=result)

                t = datetime.datetime.now()

                #print "Start time: " + str(datetime.datetime.now())

		# TODO: Move freqUQ out of this function if possible.
                class freqUQ:
                    def __init__(self):
                        self.list=[]
                    def step(self, value):
                        self.list.append(value)
                    def finalize(self):
                        return sorted([(self.list.count(x), x) for x in set(self.list)], reverse=True)[0][1]

                #MAKE THIS SO IT CAN ATTACH SQLITE3 FOR BENCHMARK
		# UniqueID,Patent,Lastname,Firstname
		# ,VARCHAR,,
		# ,%08d,,
		# UNIQUE,EXACT,FUZZY,FUZZY
		# 1,5773227,ALLBRITTON,NANCY L
		# TODO: Refactor into read_benchmark_csvfile()
		# This will help for unit test and provide self-documentation
                dataS = uniVert([x for x in csv.reader(open(fileS, "rb"))])
                #print dataS

		# TODO: Refactor into create_column_types()
                #1 = Variables, 2 = Type, 3 = Format (If necessary), 4 = Matching Type
                tList = ["%s %s" % (dataS[0][i], x) for i,x in enumerate(dataS[1]) if  x != ""]

		# Slice out rows 1, 2 & 3 from dataS. This is the data which gets
		# put in the database created below, and used for matching.
		# TODO: Refactor into slice_header_from()
                dataS2 = [dataS[0]]
                dataS2.extend(dataS[4:])

                # Format if its necessary --> Basically for Patents..
                for i,x in enumerate(dataS[2]):
                    if x!="":
                        for j in xrange(1,len(dataS2)):
                            if dataS2[j][i].isdigit():
                                dataS2[j][i] = x % int(dataS2[j][i])

                conn = sqlite3.connect(":memory:")
                #conn = sqlite3.connect("bmark2.sqlite3")
                conn.create_function("jarow", 2, jarow)
                conn.create_function("errD", 2, lambda x,y: (x!=y) and 1 or None)
                conn.create_aggregate("freqUQ", 1, freqUQ)
                c = conn.cursor()

                # FIGURE OUT WHICH ONES HAVE EXACT/FUZZY
                exact = [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="E"]
                fuzzy = [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="F"]
                uqS =   [dataS[0][i] for i,x in enumerate(dataS[3]) if x.upper()[0]=="U"][0]
                #print_match_qualifiers(exact, fuzzy, uqS)

                # CREATE INDEX, MERGE DATA BASED ON EXACTS
                #print "Creating indices... " + str(datetime.datetime.now())

                exAnd = " AND ".join(["a.%s=b.%s" % (x, x) for x in exact])
		#print "exAnd: ", exAnd
                exCom = ", ".join(exact)
                #print "exCom: ", exCom

                fBnme = attach_database(c, fileB, tblB, exCom, exAnd)

                #print_diagnostics(dataS2, "dataS", True, tList)
                #quickSQL(c, data=dataS2, table="dataS", header=True, typeList=tList)
                quickSQL2(c, data=dataS2, table="dataS", header=True, typeList=tList)

		handle_dataS(c, exCom, uqB, uqS, fuzzy, fBnme, exAnd)

                create_match_tables(c, fBnme, uqB, exCom, exAnd)

                #print "Indices Done ... " + str(datetime.datetime.now())

		export_csv_results(c, output)
		print_results(c, output, t)
                c.close()
                conn.close()

if __name__ == "__main__":
    import sys
    if(sys.argv[1] == 'help' or sys.argv[1] == '?'):
        print bmVerify.__doc__

    else:
            bmVerify(sys.argv[3:], sys.argv[1], sys.argv[2])
<|endoftext|>


<commit_before>#!/usr/bin/python
# coding:utf-8

import user_manage
from pymongo import MongoClient

"""
class base_info():
    pass

class channel_info(base_info):
    def __init__(self, config=None):
        logger.debug("")
        # ch視聴情報を保存する
        #   key  : (マルチキャストアドレス, サーバのIP)
        #   value: {データパスID: channel_switch_info}
        self.channel_info = {}

        # channel_user_infoをtimeの昇順に保持するタイムアウト判定用リスト
        self.user_info_list = []

        # DBアクセサクラスのインスタンス生成
        connect_str = config.get(DB_CONNECT_STR)
        self.accessor = DatabaseAccessor(connect_str)

    def dump_self(self):
        return cPickle.dumps(self.channel_info)

    def update_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 対象ユーザーが存在しない場合は追加、存在する場合は更新処理を呼び出す
        logger.debug("")

        user = self.exists_user(mc_addr, serv_ip, datapathid, port_no, cid)
        ret = None
        if not user:
            # 追加処理
            ret = self.add_ch_info(
                mc_addr, serv_ip, datapathid, port_no, cid)
        else:
            # 更新処理
            self.update_user_info(user)
            ret = mld_const.CON_REPLY_NOTHING
        # DBへ投入
        self.accessor.upsert("viewerdata", self)
        return ret

    def add_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 視聴端末情報を追加
        logger.debug("")

        # チャンネル存在チェック
        if (mc_addr, serv_ip) not in self.channel_info:
            # 当該チャンネルが存在しない場合
            sw_info = channel_switch_info()
            sw_info.add_sw_info(mc_addr, serv_ip, datapathid, port_no, cid)
            self.channel_info[(mc_addr, serv_ip)] = {datapathid: sw_info}
            self.user_info_list.append(sw_info.port_info[port_no][cid])
            # エッジSW、収容SW両方へのFlowMod、およびエッジルータへのReport
            return mld_const.CON_REPLY_ADD_MC_GROUP

        # 収容SW(detapathid)存在チェック
        sw_info = self.channel_info[(mc_addr, serv_ip)]
        if datapathid not in sw_info:
            # 存在しない場合、既存チャンネル配下に新規収容SWの追加
            new_sw_info = channel_switch_info()
            new_sw_info.add_sw_info(
                mc_addr, serv_ip, datapathid, port_no, cid)
            sw_info[datapathid] = new_sw_info
            self.user_info_list.append(
                sw_info[datapathid].port_info[port_no][cid])
            # 収容SWへのFlowMod
            return mld_const.CON_REPLY_ADD_SWITCH

        # 当該チャンネルにこの収容SWの情報がある場合
        #   ポート情報及びユーザー情報を追加
        ch_sw_info = sw_info[datapathid]
        ret = ch_sw_info.add_sw_info(
            mc_addr, serv_ip, datapathid, port_no, cid)
        self.user_info_list.append(ch_sw_info.port_info[port_no][cid])
        return ret

    def update_user_info(self, ch_user_info):
        # 引数のchannel_user_infoのtimeを更新し、user_info_listに入れ直す
        logger.debug("")

        self.user_info_list.pop(self.user_info_list.index(ch_user_info))
        ch_user_info.time = time.time()
        self.user_info_list.append(ch_user_info)
        logger.debug("updated user_info")

    def remove_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 視聴端末を削除
        logger.debug("")

        # 削除対象の存在チェック
        user = self.exists_user(mc_addr, serv_ip, datapathid, port_no, cid)
        if not user:
            # 存在しなければ何もしない
            logger.debug("remove target is nothing.")
            return None

        # ポート以下の情報を削除
        ch_sw_info = self.channel_info[(mc_addr, serv_ip)][datapathid]
        ret = ch_sw_info.remove_sw_info(port_no, cid)

        # ポートの削除を行い、ポート配下の視聴ユーザが0になった場合
        if ret == mld_const.CON_REPLY_DEL_PORT \
                and len(ch_sw_info.port_info) == 0:

            # 収容SW(datapathid)に対応する情報を削除する
            self.channel_info[(mc_addr, serv_ip)].pop(datapathid)
            logger.debug("removed datapath : %s",  datapathid)
            ret = mld_const.CON_REPLY_DEL_SWITCH

            # 当該mcグループの視聴ユーザが0になった場合、mcグループに対応する情報を削除する
            if len(self.channel_info[(mc_addr, serv_ip)]) == 0:
                self.channel_info.pop((mc_addr, serv_ip))
                ret = mld_const.CON_REPLY_DEL_MC_GROUP

        # user_info_listから対象ユーザーを削除する
        self.user_info_list.pop(self.user_info_list.index(user))
        # DBへ投入
        self.accessor.upsert("viewerdata", self)
        return ret

    def exists_user(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # channel_infoから指定されたcidまでを持つchannel_user_infoを返却
        # 存在しない場合はNoneを返却
        logger.debug("")

        if (mc_addr, serv_ip) in self.channel_info:
            logger.debug("hit mc_addr, serv_ip : %s, %s", mc_addr, serv_ip)
            sw_info = self.channel_info[(mc_addr, serv_ip)]
            if datapathid in sw_info:
                logger.debug("hit datapathid : %s", datapathid)
                sw_info = sw_info[datapathid]
                if port_no in sw_info.port_info:
                    logger.debug("hit port_no : %s", port_no)
                    user_info = sw_info.port_info[port_no]
                    if cid in user_info:
                        logger.debug("hit cid : %s", cid)
                        logger.debug("user exists")
                        return user_info[cid]

        logger.debug("user does not exist")
        return None

    def find_insert_point(self, target_user):
        # timeでソート済みのユーザーリストから対象ユーザーを挿入すべき場所を返却
        logger.debug("")

        idx = bisect.bisect_left(self.user_info_list, target_user)
        logger.debug("idx : %d", idx)
        return idx

    def get_channel_info(self):
        # channel_infoの内容をStringで返却。（DEBUG用）

        info = "{\n"
        for key in self.channel_info.keys():
            info += "  multicast address : (%s, %s)\n" % (key[0], key[1])
            info += "  switches : [\n"
            switch_info = self.channel_info[key]
            for datapath in switch_info.keys():
                info += "    datapathid : %s\n" % datapath
                info += "    ports : [\n"
                sw_info = switch_info[datapath]
                info += sw_info.get_switch_info()
            info += "  ]\n"
        info += "}"
        return info

    def get_user_info_list(self):
        # user_info_listの内容をStringで返却。（DEBUG用）

        info = "\n"
        for user in self.user_info_list:
            info += user.get_user_info()
        return info

class channel_switch_info(base_info):
    def __init__(self):
        logger.debug("")
        # SW配下のポートの視聴情報を保存する
        #   key  : port_no
        #   value: {cid: channel_user_info}
        self.port_info = {}

    def add_sw_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # port_infoにユーザ情報を追加
        logger.debug("")
        logger.debug("port_no, cid : %s, %s", str(port_no), str(cid))

        if port_no not in self.port_info:
            # 当該ポートに視聴ユーザが存在しない場合
            ch_user_info = channel_user_info(
                mc_addr, serv_ip, datapathid, port_no, cid, time.time())
            self.port_info[port_no] = {cid: ch_user_info}
            # 収容SWへのFlowMod
            return mld_const.CON_REPLY_ADD_PORT
        else:
            # 当該ポートに視聴ユーザが存在する場合
            # 当該CIDが存在しない場合はCIDを追加
            new_user_info = channel_user_info(
                mc_addr, serv_ip, datapathid, port_no, cid, time.time())
            self.port_info[port_no][cid] = new_user_info

            # FlowMod必要なし
            return mld_const.CON_REPLY_NOTHING

    def remove_sw_info(self, port_no, cid):
        # port_infoからユーザ情報を削除
        logger.debug("")

        # port_infoから当該ユーザ情報を検索し削除
        self.port_info[port_no].pop(cid)
        logger.debug("removed user[cid : %s]", cid)
        if len(self.port_info[port_no]) == 0:
            # ポート配下のユーザが0になればポート情報も削除
            self.port_info.pop(port_no)
            # 収容SWへのFlowModが必要
            return mld_const.CON_REPLY_DEL_PORT
        else:
            # FlowMod必要なし
            return mld_const.CON_REPLY_NOTHING

    def get_switch_info(self):
        # port_infoの内容をStringで返却。（DEBUG用）

        info = ""
        for port in self.port_info.keys():
            info += "      port : %s\n" % port
            info += "      users : [\n"
            user_info = self.port_info[port]
            for cid in user_info.keys():
                ch_user_info = user_info[cid]
                info += ch_user_info.get_user_info()
        info += "    ]\n"
        return info

@total_ordering
class channel_user_info(base_info):
    def __init__(self, mc_addr, serv_ip, datapathid, port_no, cid, time):
        logger.debug("")
        self.mc_addr = mc_addr
        self.serv_ip = serv_ip
        self.datapathid = datapathid
        self.port_no = port_no
        self.cid = cid
        self.time = time

    def get_user_info(self):
        # user_infoの内容をStringで返却。（DEBUG用）
        info = ""
        info += "        {\n"
#        info += "          mc_addr : %s\n" % self.mc_addr
#        info += "          serv_ip : %s\n" % self.serv_ip
#        info += "          datapathid : %s\n" % self.datapathid
#        info += "          port_no : %s\n" % self.port_no
        info += "          cid : %s\n" % self.cid
        info += "          time : %f\n" % self.time
        info += "        }\n"
        return info

    def __eq__(self, other):
        return (self.time == other.time)

    def __lt__(self, other):
        return (self.time < other.time)

# 動作確認用
# if __name__ == "__main__":
#    ch_info = channel_info()
#
#    # add
#    logger.debug("**** <1>")
#    ch_info.update_ch_info("ff38::1:1", "2001::1:20", 1, 1, 111)
#    logger.debug("")
#    logger.debug("**** <1> add cid 100")
#    ch_info.update_ch_info("ff38::1:1", "2001::1:20", 1, 1, 100)
#    logger.debug("")
#    logger.debug("**** <1> add port 2")
#    ch_info.update_ch_info("ff38::1:1", "2001::1:20", 1, 2, 120)
#    logger.debug("")
#    logger.debug("**** <1> add datapath 2")
#    ch_info.update_ch_info("ff38::1:1", "2001::1:20", 2, 1, 210)
#    logger.debug("")
#    logger.debug("**** <2>")
#    ch_info.update_ch_info("ff38::1:2", "2001::1:20", 1, 1, 111)
#    logger.debug(ch_info.get_channel_info())
#    logger.debug(ch_info.get_user_info_list())
#
#    # update
#    logger.debug("")
#    logger.debug("**** <1> update cid 111")
#    ch_info.update_ch_info("ff38::1:1", "2001::1:20", 1, 1, 111)
#    logger.debug("")
#    logger.debug(ch_info.get_channel_info())
#    logger.debug(ch_info.get_user_info_list())
#
#    # remove
#    logger.debug("")
#    logger.debug("**** remove <1> cid 111")
#    ch_info.remove_ch_info("ff38::1:1", "2001::1:20", 1, 1, 111)
#    logger.debug("")
#    ch_info.remove_ch_info("ff38::1:1", "2001::1:20", 1, 1, 100)
#    logger.debug("")
#    logger.debug("**** remove <1> datapath 2")
#    ch_info.remove_ch_info("ff38::1:1", "2001::1:20", 2, 1, 210)
#    logger.debug("")
#    logger.debug("**** remove <1> port 2")
#    ch_info.remove_ch_info("ff38::1:1", "2001::1:20", 1, 2, 120)
#    logger.debug("")
#    logger.debug("**** remove <2>")
#    ch_info.remove_ch_info("ff38::1:2", "2001::1:20", 1, 1, 111)
#    logger.debug(ch_info.get_channel_info())
#    logger.debug(ch_info.get_user_info_list())

class DatabaseAccessor:
    def __init__(self, connect_str):
        self.client = None
        if not connect_str:
            return
        self.client = MongoClient(connect_str)
        # DB名:viewerdb、コレクション名:serialized_data
        self.db = self.client.viewerdb
        self.col = self.db.serialized_data

    def upsert(self, key, inserted_obj):
        if not self.client:
            return
        # 投入対象オブジェクトをdumpしてそのまま投入
        dump = inserted_obj.dump_self()
        self.col.update({"ch": "all"}, {"$set": {key: dump}}, upsert=True)

    def query(self, key):
        if not self.client:
            return None
        result = self.col.find_one()
        dump = result[key]
        return cPickle.loads(dump)

"""

if '__main__' == __name__:
    # open mongodb
    client = MongoClient("mongodb://localhost:27017")
    db = client.viewerdb
    col = db.serialized_data

    # query
    result = col.find_one({"ch": "all"})
    dump_result = result["viewerdata"]
    load_result = cpickle.loads(dump_result)

    # check
    logger.debug(load_result)
<commit_msg>ユーザ管理DB表示ツール(仮)の表示を改善<commit_after>#!/usr/bin/python
# coding:utf-8

import user_manage
import cPickle
import time
from pymongo import MongoClient

"""
class base_info():
    pass

class channel_info(base_info):
    def __init__(self, config=None):
        logger.debug("")
        # ch視聴情報を保存する
        #   key  : (マルチキャストアドレス, サーバのIP)
        #   value: {データパスID: channel_switch_info}
        self.channel_info = {}

        # channel_user_infoをtimeの昇順に保持するタイムアウト判定用リスト
        self.user_info_list = []

        # DBアクセサクラスのインスタンス生成
        connect_str = config.get(DB_CONNECT_STR)
        self.accessor = DatabaseAccessor(connect_str)

    def dump_self(self):
        return cPickle.dumps(self.channel_info)

    def update_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 対象ユーザーが存在しない場合は追加、存在する場合は更新処理を呼び出す
        logger.debug("")

        user = self.exists_user(mc_addr, serv_ip, datapathid, port_no, cid)
        ret = None
        if not user:
            # 追加処理
            ret = self.add_ch_info(
                mc_addr, serv_ip, datapathid, port_no, cid)
        else:
            # 更新処理
            self.update_user_info(user)
            ret = mld_const.CON_REPLY_NOTHING
        # DBへ投入
        self.accessor.upsert("viewerdata", self)
        return ret

    def add_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 視聴端末情報を追加
        logger.debug("")

        # チャンネル存在チェック
        if (mc_addr, serv_ip) not in self.channel_info:
            # 当該チャンネルが存在しない場合
            sw_info = channel_switch_info()
            sw_info.add_sw_info(mc_addr, serv_ip, datapathid, port_no, cid)
            self.channel_info[(mc_addr, serv_ip)] = {datapathid: sw_info}
            self.user_info_list.append(sw_info.port_info[port_no][cid])
            # エッジSW、収容SW両方へのFlowMod、およびエッジルータへのReport
            return mld_const.CON_REPLY_ADD_MC_GROUP

        # 収容SW(detapathid)存在チェック
        sw_info = self.channel_info[(mc_addr, serv_ip)]
        if datapathid not in sw_info:
            # 存在しない場合、既存チャンネル配下に新規収容SWの追加
            new_sw_info = channel_switch_info()
            new_sw_info.add_sw_info(
                mc_addr, serv_ip, datapathid, port_no, cid)
            sw_info[datapathid] = new_sw_info
            self.user_info_list.append(
                sw_info[datapathid].port_info[port_no][cid])
            # 収容SWへのFlowMod
            return mld_const.CON_REPLY_ADD_SWITCH

        # 当該チャンネルにこの収容SWの情報がある場合
        #   ポート情報及びユーザー情報を追加
        ch_sw_info = sw_info[datapathid]
        ret = ch_sw_info.add_sw_info(
            mc_addr, serv_ip, datapathid, port_no, cid)
        self.user_info_list.append(ch_sw_info.port_info[port_no][cid])
        return ret

    def update_user_info(self, ch_user_info):
        # 引数のchannel_user_infoのtimeを更新し、user_info_listに入れ直す
        logger.debug("")

        self.user_info_list.pop(self.user_info_list.index(ch_user_info))
        ch_user_info.time = time.time()
        self.user_info_list.append(ch_user_info)
        logger.debug("updated user_info")

    def remove_ch_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # 視聴端末を削除
        logger.debug("")

        # 削除対象の存在チェック
        user = self.exists_user(mc_addr, serv_ip, datapathid, port_no, cid)
        if not user:
            # 存在しなければ何もしない
            logger.debug("remove target is nothing.")
            return None

        # ポート以下の情報を削除
        ch_sw_info = self.channel_info[(mc_addr, serv_ip)][datapathid]
        ret = ch_sw_info.remove_sw_info(port_no, cid)

        # ポートの削除を行い、ポート配下の視聴ユーザが0になった場合
        if ret == mld_const.CON_REPLY_DEL_PORT \
                and len(ch_sw_info.port_info) == 0:

            # 収容SW(datapathid)に対応する情報を削除する
            self.channel_info[(mc_addr, serv_ip)].pop(datapathid)
            logger.debug("removed datapath : %s",  datapathid)
            ret = mld_const.CON_REPLY_DEL_SWITCH

            # 当該mcグループの視聴ユーザが0になった場合、mcグループに対応する情報を削除する
            if len(self.channel_info[(mc_addr, serv_ip)]) == 0:
                self.channel_info.pop((mc_addr, serv_ip))
                ret = mld_const.CON_REPLY_DEL_MC_GROUP

        # user_info_listから対象ユーザーを削除する
        self.user_info_list.pop(self.user_info_list.index(user))
        # DBへ投入
        self.accessor.upsert("viewerdata", self)
        return ret

    def exists_user(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # channel_infoから指定されたcidまでを持つchannel_user_infoを返却
        # 存在しない場合はNoneを返却
        logger.debug("")

        if (mc_addr, serv_ip) in self.channel_info:
            logger.debug("hit mc_addr, serv_ip : %s, %s", mc_addr, serv_ip)
            sw_info = self.channel_info[(mc_addr, serv_ip)]
            if datapathid in sw_info:
                logger.debug("hit datapathid : %s", datapathid)
                sw_info = sw_info[datapathid]
                if port_no in sw_info.port_info:
                    logger.debug("hit port_no : %s", port_no)
                    user_info = sw_info.port_info[port_no]
                    if cid in user_info:
                        logger.debug("hit cid : %s", cid)
                        logger.debug("user exists")
                        return user_info[cid]

        logger.debug("user does not exist")
        return None

    def find_insert_point(self, target_user):
        # timeでソート済みのユーザーリストから対象ユーザーを挿入すべき場所を返却
        logger.debug("")

        idx = bisect.bisect_left(self.user_info_list, target_user)
        logger.debug("idx : %d", idx)
        return idx

    def get_channel_info(self):
        # channel_infoの内容をStringで返却。（DEBUG用）

        info = "{\n"
        for key in self.channel_info.keys():
            info += "  multicast address : (%s, %s)\n" % (key[0], key[1])
            info += "  switches : [\n"
            switch_info = self.channel_info[key]
            for datapath in switch_info.keys():
                info += "    datapathid : %s\n" % datapath
                info += "    ports : [\n"
                sw_info = switch_info[datapath]
                info += sw_info.get_switch_info()
            info += "  ]\n"
        info += "}"
        return info

    def get_user_info_list(self):
        # user_info_listの内容をStringで返却。（DEBUG用）

        info = "\n"
        for user in self.user_info_list:
            info += user.get_user_info()
        return info

class channel_switch_info(base_info):
    def __init__(self):
        logger.debug("")
        # SW配下のポートの視聴情報を保存する
        #   key  : port_no
        #   value: {cid: channel_user_info}
        self.port_info = {}

    def add_sw_info(self, mc_addr, serv_ip, datapathid, port_no, cid):
        # port_infoにユーザ情報を追加
        logger.debug("")
        logger.debug("port_no, cid : %s, %s", str(port_no), str(cid))

        if port_no not in self.port_info:
            # 当該ポートに視聴ユーザが存在しない場合
            ch_user_info = channel_user_info(
                mc_addr, serv_ip, datapathid, port_no, cid, time.time())
            self.port_info[port_no] = {cid: ch_user_info}
            # 収容SWへのFlowMod
            return mld_const.CON_REPLY_ADD_PORT
        else:
            # 当該ポートに視聴ユーザが存在する場合
            # 当該CIDが存在しない場合はCIDを追加
            new_user_info = channel_user_info(
                mc_addr, serv_ip, datapathid, port_no, cid, time.time())
            self.port_info[port_no][cid] = new_user_info

            # FlowMod必要なし
            return mld_const.CON_REPLY_NOTHING

    def remove_sw_info(self, port_no, cid):
        # port_infoからユーザ情報を削除
        logger.debug("")

        # port_infoから当該ユーザ情報を検索し削除
        self.port_info[port_no].pop(cid)
        logger.debug("removed user[cid : %s]", cid)
        if len(self.port_info[port_no]) == 0:
            # ポート配下のユーザが0になればポート情報も削除
            self.port_info.pop(port_no)
            # 収容SWへのFlowModが必要
            return mld_const.CON_REPLY_DEL_PORT
        else:
            # FlowMod必要なし
            return mld_const.CON_REPLY_NOTHING

    def get_switch_info(self):
        # port_infoの内容をStringで返却。（DEBUG用）

        info = ""
        for port in self.port_info.keys():
            info += "      port : %s\n" % port
            info += "      users : [\n"
            user_info = self.port_info[port]
            for cid in user_info.keys():
                ch_user_info = user_info[cid]
                info += ch_user_info.get_user_info()
        info += "    ]\n"
        return info

@total_ordering
class channel_user_info(base_info):
    def __init__(self, mc_addr, serv_ip, datapathid, port_no, cid, time):
        logger.debug("")
        self.mc_addr = mc_addr
        self.serv_ip = serv_ip
        self.datapathid = datapathid
        self.port_no = port_no
        self.cid = cid
        self.time = time

    def get_user_info(self):
        # user_infoの内容をStringで返却。（DEBUG用）
        info = ""
        info += "        {\n"
#        info += "          mc_addr : %s\n" % self.mc_addr
#        info += "          serv_ip : %s\n" % self.serv_ip
#        info += "          datapathid : %s\n" % self.datapathid
#        info += "          port_no : %s\n" % self.port_no
        info += "          cid : %s\n" % self.cid
        info += "          time : %f\n" % self.time
        info += "        }\n"
        return info

    def __eq__(self, other):
        return (self.time == other.time)

    def __lt__(self, other):
        return (self.time < other.time)

class DatabaseAccessor:
    def __init__(self, connect_str):
        self.client = None
        if not connect_str:
            return
        self.client = MongoClient(connect_str)
        # DB名:viewerdb、コレクション名:serialized_data
        self.db = self.client.viewerdb
        self.col = self.db.serialized_data

    def upsert(self, key, inserted_obj):
        if not self.client:
            return
        # 投入対象オブジェクトをdumpしてそのまま投入
        dump = inserted_obj.dump_self()
        self.col.update({"ch": "all"}, {"$set": {key: dump}}, upsert=True)

    def query(self, key):
        if not self.client:
            return None
        result = self.col.find_one()
        dump = result[key]
        return cPickle.loads(dump)

"""

if '__main__' == __name__:
    # open mongodb
    client = MongoClient("mongodb://localhost:27017")
    db = client.viewerdb
    col = db.serialized_data

    # query
    result = col.find_one({"ch": "all"})
    dump_result = result["viewerdata"]
    load_result = cPickle.loads(str(dump_result))

    # display
    for (key, switch_dict) in load_result.items():
        print key
        for (datapath_id, sw_obj) in switch_dict.items():
            print "\tSwitch:", datapath_id
            for (port_no, usr_dict) in sw_obj.port_info.items():
                print "\t\tPort_No:", port_no
                for (cid, usr_obj) in usr_dict.items():
                    print "\t\t\tcid: ", cid, " Update time: ", time.strftime("%Y/%m/%d %H:%M:%S", time.localtime(usr_obj.time))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Copyright(C) 2010-2014 Romain Bignon
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

import os

from weboob.core.bcall import BackendsCall
from weboob.core.modules import ModulesLoader, RepositoryModulesLoader, ModuleLoadError
from weboob.core.backendscfg import BackendsConfig
from weboob.core.repositories import Repositories, PrintProgress
from weboob.core.scheduler import Scheduler
from weboob.tools.backend import Module
from weboob.tools.config.iconfig import ConfigError
from weboob.tools.log import getLogger

__all__ = ['WebNip', 'Weboob']

class VersionsMismatchError(ConfigError):
    pass

class WebNip(object):
    """
    Weboob in Non Integrated Programs

    It provides methods to build backends or call methods on all loaded
    backends.

    :param modules_path: path to directory containing modules.
    :type modules_path: :class:`basestring`
    :param storage: provide a storage where backends can save data
    :type storage: :class:`weboob.tools.storage.IStorage`
    :param scheduler: what scheduler to use; default is :class:`weboob.core.scheduler.Scheduler`
    :type scheduler: :class:`weboob.core.scheduler.IScheduler`
    """
    VERSION = '1.1'

    def __init__(self, modules_path=None, storage=None, scheduler=None):
        self.logger = getLogger('weboob')
        self.backend_instances = {}
        self.callbacks = {'login':   lambda backend_name, value: None,
                          'captcha': lambda backend_name, image: None,
                         }

        if modules_path is None:
            import pkg_resources
            modules_path = pkg_resources.resource_filename('weboob_modules', '')

        if modules_path:
            self.modules_loader = ModulesLoader(modules_path, self.VERSION)

        if scheduler is None:
            scheduler = Scheduler()
        self.scheduler = scheduler

        self.storage = storage

    def __deinit__(self):
        self.deinit()

    def deinit(self):
        """
        Call this method when you stop using Weboob, to
        properly unload all correctly.
        """
        self.unload_backends()

    def build_backend(self, module_name, params=None, storage=None, name=None):
        """
        Create a backend.

        It does not load it into the Weboob object, so you are responsible for
        deinitialization and calls.

        :param module_name: name of module
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param name: name of backend
        :type name: :class:`basestring`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        module = self.modules_loader.get_or_load_module(module_name)

        backend_instance = module.create_instance(self, name or module_name, params or {}, storage)
        return backend_instance

    class LoadError(Exception):
        """
        Raised when a backend is unabled to load.

        :param backend_name: name of backend we can't load
        :param exception: exception object
        """

        def __init__(self, backend_name, exception):
            Exception.__init__(self, unicode(exception))
            self.backend_name = backend_name

    def load_backend(self, module_name, name, params=None, storage=None):
        """
        Load a backend.

        :param module_name: name of module to load
        :type module_name: :class:`basestring`:
        :param name: name of instance
        :type name: :class:`basestring`
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        if name is None:
            name = module_name

        if name in self.backend_instances:
            raise self.LoadError(name, 'A loaded backend already named "%s"' % name)

        backend = self.build_backend(module_name, params, storage, name)
        self.backend_instances[name] = backend
        return backend

    def unload_backends(self, names=None):
        """
        Unload backends.

        :param names: if specified, only unload that backends
        :type names: :class:`list`
        """
        unloaded = {}
        if isinstance(names, basestring):
            names = [names]
        elif names is None:
            names = self.backend_instances.keys()

        for name in names:
            backend = self.backend_instances.pop(name)
            with backend:
                backend.deinit()
            unloaded[backend.name] = backend

        return unloaded

    def __getitem__(self, name):
        """
        Alias for :func:`WebNip.get_backend`.
        """
        return self.get_backend(name)

    def get_backend(self, name, **kwargs):
        """
        Get a backend from its name.

        :param name: name of backend to get
        :type name: str
        :param default: if specified, get this value when the backend is not found
        :type default: whatever you want
        :raises: :class:`KeyError` if not found.
        """
        try:
            return self.backend_instances[name]
        except KeyError:
            if 'default' in kwargs:
                return kwargs['default']
            else:
                raise

    def count_backends(self):
        """
        Get number of loaded backends.
        """
        return len(self.backend_instances)

    def iter_backends(self, caps=None, module=None):
        """
        Iter on each backends.

        Note: each backend is locked when it is returned.

        :param caps: optional list of capabilities to select backends
        :type caps: tuple[:class:`weboob.capabilities.base.Capability`]
        :param module: optional name of module
        :type module: :class:`basestring`
        :rtype: iter[:class:`weboob.tools.backend.Module`]
        """
        for _, backend in sorted(self.backend_instances.iteritems()):
            if (caps is None or backend.has_caps(caps)) and \
               (module is None or backend.NAME == module):
                with backend:
                    yield backend

    def __getattr__(self, name):
        def caller(*args, **kwargs):
            return self.do(name, *args, **kwargs)
        return caller

    def do(self, function, *args, **kwargs):
        r"""
        Do calls on loaded backends with specified arguments, in separated
        threads.

        This function has two modes:

        - If *function* is a string, it calls the method with this name on
          each backends with the specified arguments;
        - If *function* is a callable, it calls it in a separated thread with
          the locked backend instance at first arguments, and \*args and
          \*\*kwargs.

        :param function: backend's method name, or a callable object
        :type function: :class:`str`
        :param backends: list of backends to iterate on
        :type backends: list[:class:`str`]
        :param caps: iterate on backends which implement this caps
        :type caps: list[:class:`weboob.capabilities.base.Capability`]
        :rtype: A :class:`weboob.core.bcall.BackendsCall` object (iterable)
        """
        backends = self.backend_instances.values()
        _backends = kwargs.pop('backends', None)
        if _backends is not None:
            if isinstance(_backends, Module):
                backends = [_backends]
            elif isinstance(_backends, basestring):
                if len(_backends) > 0:
                    try:
                        backends = [self.backend_instances[_backends]]
                    except (ValueError, KeyError):
                        backends = []
            elif isinstance(_backends, (list, tuple, set)):
                backends = []
                for backend in _backends:
                    if isinstance(backend, basestring):
                        try:
                            backends.append(self.backend_instances[backend])
                        except (ValueError, KeyError):
                            pass
                    else:
                        backends.append(backend)
            else:
                self.logger.warning(u'The "backends" value isn\'t supported: %r', _backends)

        if 'caps' in kwargs:
            caps = kwargs.pop('caps')
            backends = [backend for backend in backends if backend.has_caps(caps)]

        # The return value MUST BE the BackendsCall instance. Please never iterate
        # here on this object, because caller might want to use other methods, like
        # wait() on callback_thread().
        # Thanks a lot.
        return BackendsCall(backends, function, *args, **kwargs)

    def schedule(self, interval, function, *args):
        """
        Schedule an event.

        :param interval: delay before calling the function
        :type interval: int
        :param function: function to call
        :type function: callabale
        :param args: arguments to give to function
        :returns: an event identificator
        """
        return self.scheduler.schedule(interval, function, *args)

    def repeat(self, interval, function, *args):
        """
        Repeat a call to a function

        :param interval: interval between two calls
        :type interval: int
        :param function: function to call
        :type function: callable
        :param args: arguments to give to function
        :returns: an event identificator
        """
        return self.scheduler.repeat(interval, function, *args)

    def cancel(self, ev):
        """
        Cancel an event

        :param ev: the event identificator
        """
        return self.scheduler.cancel(ev)

    def want_stop(self):
        """
        Plan to stop the scheduler.
        """
        return self.scheduler.want_stop()

    def loop(self):
        """
        Run the scheduler loop
        """
        return self.scheduler.run()

class Weboob(WebNip):
    """
    The main class of Weboob, used to manage backends, modules repositories and
    call methods on all loaded backends.

    :param workdir: optional parameter to set path of the working directory
    :type workdir: str
    :param backends_filename: name of the *backends* file, where configuration of
                              backends is stored
    :type backends_filename: str
    :param storage: provide a storage where backends can save data
    :type storage: :class:`weboob.tools.storage.IStorage`
    """
    BACKENDS_FILENAME = 'backends'

    def __init__(self, workdir=None, backends_filename=None, scheduler=None, storage=None):
        super(Weboob, self).__init__(modules_path=False, scheduler=scheduler, storage=storage)

        # Create WORKDIR
        if workdir is not None:
            datadir = workdir
        elif 'WEBOOB_WORKDIR' in os.environ:
            datadir = workdir = os.environ.get('WEBOOB_WORKDIR')
        else:
            workdir = os.path.join(os.environ.get('XDG_CONFIG_HOME', os.path.join(os.path.expanduser('~'), '.config')), 'weboob')
            datadir = os.path.join(os.environ.get('XDG_DATA_HOME', os.path.join(os.path.expanduser('~'), '.local', 'share')), 'weboob')

        self.workdir = os.path.realpath(workdir)
        self._create_dir(workdir)

        # Modules management
        self.repositories = Repositories(workdir, datadir, self.VERSION)
        self.modules_loader = RepositoryModulesLoader(self.repositories)

        # Backend instances config
        if not backends_filename:
            backends_filename = os.environ.get('WEBOOB_BACKENDS', os.path.join(self.workdir, self.BACKENDS_FILENAME))
        elif not backends_filename.startswith('/'):
            backends_filename = os.path.join(self.workdir, backends_filename)
        self.backends_config = BackendsConfig(backends_filename)

    def _create_dir(self, name):
        if not os.path.exists(name):
            os.makedirs(name)
        elif not os.path.isdir(name):
            self.logger.error(u'"%s" is not a directory', name)

    def update(self, progress=PrintProgress()):
        """
        Update modules from repositories.
        """
        self.repositories.update(progress)

        modules_to_check = set([module_name for _, module_name, _ in self.backends_config.iter_backends()])
        for module_name in modules_to_check:
            minfo = self.repositories.get_module_info(module_name)
            if minfo and not minfo.is_installed():
                self.repositories.install(minfo, progress)

    def build_backend(self, module_name, params=None, storage=None, name=None):
        """
        Create a single backend which is not listed in configuration.

        :param module_name: name of module
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param name: name of backend
        :type name: :class:`basestring`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        minfo = self.repositories.get_module_info(module_name)
        if minfo is None:
            raise ModuleLoadError(module_name, 'Module does not exist.')

        if not minfo.is_installed():
            self.repositories.install(minfo)

        return super(Weboob, self).build_backend(module_name, params, storage, name)

    def load_backends(self, caps=None, names=None, modules=None, exclude=None, storage=None, errors=None):
        """
        Load backends listed in config file.

        :param caps: load backends which implement all of specified caps
        :type caps: tuple[:class:`weboob.capabilities.base.Capability`]
        :param names: load backends with instance name in list
        :type names: tuple[:class:`str`]
        :param modules: load backends which module is in list
        :type modules: tuple[:class:`str`]
        :param exclude: do not load modules in list
        :type exclude: tuple[:class:`str`]
        :param storage: use this storage if specified
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param errors: if specified, store every errors in this list
        :type errors: list[:class:`LoadError`]
        :returns: loaded backends
        :rtype: dict[:class:`str`, :class:`weboob.tools.backend.Module`]
        """
        loaded = {}
        if storage is None:
            storage = self.storage

        if not self.repositories.check_repositories():
            self.logger.error(u'Repositories are not consistent with the sources.list')
            raise VersionsMismatchError(u'Versions mismatch, please run "weboob-config update"')

        for instance_name, module_name, params in self.backends_config.iter_backends():
            if '_enabled' in params and not params['_enabled'].lower() in ('1', 'y', 'true', 'on', 'yes') or \
               names is not None and instance_name not in names or \
               modules is not None and module_name not in modules or \
               exclude is not None and module_name in exclude:
                continue

            minfo = self.repositories.get_module_info(module_name)
            if minfo is None:
                self.logger.warning(u'Backend "%s" is referenced in %s but was not found. '
                                    u'Perhaps a missing repository?', module_name, self.backends_config.confpath)
                continue

            if caps is not None and not minfo.has_caps(caps):
                continue

            if not minfo.is_installed():
                self.repositories.install(minfo)

            module = None
            try:
                module = self.modules_loader.get_or_load_module(module_name)
            except ModuleLoadError as e:
                self.logger.error(u'Unable to load module "%s": %s', module_name, e)
                continue

            if instance_name in self.backend_instances:
                self.logger.warning(u'Oops, the backend "%s" is already loaded. Unload it before reloading...', instance_name)
                self.unload_backends(instance_name)

            try:
                backend_instance = module.create_instance(self, instance_name, params, storage)
            except Module.ConfigError as e:
                if errors is not None:
                    errors.append(self.LoadError(instance_name, e))
            else:
                self.backend_instances[instance_name] = loaded[instance_name] = backend_instance
        return loaded
<commit_msg>[core] add a datadir option to weboob init and handle WEBOOB_DATADIR<commit_after># -*- coding: utf-8 -*-

# Copyright(C) 2010-2014 Romain Bignon
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

import os

from weboob.core.bcall import BackendsCall
from weboob.core.modules import ModulesLoader, RepositoryModulesLoader, ModuleLoadError
from weboob.core.backendscfg import BackendsConfig
from weboob.core.repositories import Repositories, PrintProgress
from weboob.core.scheduler import Scheduler
from weboob.tools.backend import Module
from weboob.tools.config.iconfig import ConfigError
from weboob.tools.log import getLogger

__all__ = ['WebNip', 'Weboob']

class VersionsMismatchError(ConfigError):
    pass

class WebNip(object):
    """
    Weboob in Non Integrated Programs

    It provides methods to build backends or call methods on all loaded
    backends.

    :param modules_path: path to directory containing modules.
    :type modules_path: :class:`basestring`
    :param storage: provide a storage where backends can save data
    :type storage: :class:`weboob.tools.storage.IStorage`
    :param scheduler: what scheduler to use; default is :class:`weboob.core.scheduler.Scheduler`
    :type scheduler: :class:`weboob.core.scheduler.IScheduler`
    """
    VERSION = '1.1'

    def __init__(self, modules_path=None, storage=None, scheduler=None):
        self.logger = getLogger('weboob')
        self.backend_instances = {}
        self.callbacks = {'login':   lambda backend_name, value: None,
                          'captcha': lambda backend_name, image: None,
                         }

        if modules_path is None:
            import pkg_resources
            modules_path = pkg_resources.resource_filename('weboob_modules', '')

        if modules_path:
            self.modules_loader = ModulesLoader(modules_path, self.VERSION)

        if scheduler is None:
            scheduler = Scheduler()
        self.scheduler = scheduler

        self.storage = storage

    def __deinit__(self):
        self.deinit()

    def deinit(self):
        """
        Call this method when you stop using Weboob, to
        properly unload all correctly.
        """
        self.unload_backends()

    def build_backend(self, module_name, params=None, storage=None, name=None):
        """
        Create a backend.

        It does not load it into the Weboob object, so you are responsible for
        deinitialization and calls.

        :param module_name: name of module
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param name: name of backend
        :type name: :class:`basestring`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        module = self.modules_loader.get_or_load_module(module_name)

        backend_instance = module.create_instance(self, name or module_name, params or {}, storage)
        return backend_instance

    class LoadError(Exception):
        """
        Raised when a backend is unabled to load.

        :param backend_name: name of backend we can't load
        :param exception: exception object
        """

        def __init__(self, backend_name, exception):
            Exception.__init__(self, unicode(exception))
            self.backend_name = backend_name

    def load_backend(self, module_name, name, params=None, storage=None):
        """
        Load a backend.

        :param module_name: name of module to load
        :type module_name: :class:`basestring`:
        :param name: name of instance
        :type name: :class:`basestring`
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        if name is None:
            name = module_name

        if name in self.backend_instances:
            raise self.LoadError(name, 'A loaded backend already named "%s"' % name)

        backend = self.build_backend(module_name, params, storage, name)
        self.backend_instances[name] = backend
        return backend

    def unload_backends(self, names=None):
        """
        Unload backends.

        :param names: if specified, only unload that backends
        :type names: :class:`list`
        """
        unloaded = {}
        if isinstance(names, basestring):
            names = [names]
        elif names is None:
            names = self.backend_instances.keys()

        for name in names:
            backend = self.backend_instances.pop(name)
            with backend:
                backend.deinit()
            unloaded[backend.name] = backend

        return unloaded

    def __getitem__(self, name):
        """
        Alias for :func:`WebNip.get_backend`.
        """
        return self.get_backend(name)

    def get_backend(self, name, **kwargs):
        """
        Get a backend from its name.

        :param name: name of backend to get
        :type name: str
        :param default: if specified, get this value when the backend is not found
        :type default: whatever you want
        :raises: :class:`KeyError` if not found.
        """
        try:
            return self.backend_instances[name]
        except KeyError:
            if 'default' in kwargs:
                return kwargs['default']
            else:
                raise

    def count_backends(self):
        """
        Get number of loaded backends.
        """
        return len(self.backend_instances)

    def iter_backends(self, caps=None, module=None):
        """
        Iter on each backends.

        Note: each backend is locked when it is returned.

        :param caps: optional list of capabilities to select backends
        :type caps: tuple[:class:`weboob.capabilities.base.Capability`]
        :param module: optional name of module
        :type module: :class:`basestring`
        :rtype: iter[:class:`weboob.tools.backend.Module`]
        """
        for _, backend in sorted(self.backend_instances.iteritems()):
            if (caps is None or backend.has_caps(caps)) and \
               (module is None or backend.NAME == module):
                with backend:
                    yield backend

    def __getattr__(self, name):
        def caller(*args, **kwargs):
            return self.do(name, *args, **kwargs)
        return caller

    def do(self, function, *args, **kwargs):
        r"""
        Do calls on loaded backends with specified arguments, in separated
        threads.

        This function has two modes:

        - If *function* is a string, it calls the method with this name on
          each backends with the specified arguments;
        - If *function* is a callable, it calls it in a separated thread with
          the locked backend instance at first arguments, and \*args and
          \*\*kwargs.

        :param function: backend's method name, or a callable object
        :type function: :class:`str`
        :param backends: list of backends to iterate on
        :type backends: list[:class:`str`]
        :param caps: iterate on backends which implement this caps
        :type caps: list[:class:`weboob.capabilities.base.Capability`]
        :rtype: A :class:`weboob.core.bcall.BackendsCall` object (iterable)
        """
        backends = self.backend_instances.values()
        _backends = kwargs.pop('backends', None)
        if _backends is not None:
            if isinstance(_backends, Module):
                backends = [_backends]
            elif isinstance(_backends, basestring):
                if len(_backends) > 0:
                    try:
                        backends = [self.backend_instances[_backends]]
                    except (ValueError, KeyError):
                        backends = []
            elif isinstance(_backends, (list, tuple, set)):
                backends = []
                for backend in _backends:
                    if isinstance(backend, basestring):
                        try:
                            backends.append(self.backend_instances[backend])
                        except (ValueError, KeyError):
                            pass
                    else:
                        backends.append(backend)
            else:
                self.logger.warning(u'The "backends" value isn\'t supported: %r', _backends)

        if 'caps' in kwargs:
            caps = kwargs.pop('caps')
            backends = [backend for backend in backends if backend.has_caps(caps)]

        # The return value MUST BE the BackendsCall instance. Please never iterate
        # here on this object, because caller might want to use other methods, like
        # wait() on callback_thread().
        # Thanks a lot.
        return BackendsCall(backends, function, *args, **kwargs)

    def schedule(self, interval, function, *args):
        """
        Schedule an event.

        :param interval: delay before calling the function
        :type interval: int
        :param function: function to call
        :type function: callabale
        :param args: arguments to give to function
        :returns: an event identificator
        """
        return self.scheduler.schedule(interval, function, *args)

    def repeat(self, interval, function, *args):
        """
        Repeat a call to a function

        :param interval: interval between two calls
        :type interval: int
        :param function: function to call
        :type function: callable
        :param args: arguments to give to function
        :returns: an event identificator
        """
        return self.scheduler.repeat(interval, function, *args)

    def cancel(self, ev):
        """
        Cancel an event

        :param ev: the event identificator
        """
        return self.scheduler.cancel(ev)

    def want_stop(self):
        """
        Plan to stop the scheduler.
        """
        return self.scheduler.want_stop()

    def loop(self):
        """
        Run the scheduler loop
        """
        return self.scheduler.run()

class Weboob(WebNip):
    """
    The main class of Weboob, used to manage backends, modules repositories and
    call methods on all loaded backends.

    :param workdir: optional parameter to set path of the working directory
    :type workdir: str
    :param datadir: optional parameter to set path of the data directory
    :type datadir: str
    :param backends_filename: name of the *backends* file, where configuration of
                              backends is stored
    :type backends_filename: str
    :param storage: provide a storage where backends can save data
    :type storage: :class:`weboob.tools.storage.IStorage`
    """
    BACKENDS_FILENAME = 'backends'

    def __init__(self, workdir=None, datadir=None, backends_filename=None, scheduler=None, storage=None):
        super(Weboob, self).__init__(modules_path=False, scheduler=scheduler, storage=storage)

        # Create WORKDIR
        if workdir is None:
            if 'WEBOOB_WORKDIR' in os.environ:
                workdir = os.environ.get('WEBOOB_WORKDIR')
            else:
                workdir = os.path.join(os.environ.get('XDG_CONFIG_HOME', os.path.join(os.path.expanduser('~'), '.config')), 'weboob')

        self.workdir = os.path.realpath(workdir)
        self._create_dir(workdir)

        # Create DATADIR
        if datadir is None:
            if 'WEBOOB_DATADIR' in os.environ:
                datadir = os.environ.get('WEBOOB_DATADIR')
            else:
                datadir = os.path.join(os.environ.get('XDG_DATA_HOME', os.path.join(os.path.expanduser('~'), '.local', 'share')), 'weboob')

        _datadir = os.path.realpath(datadir)
        self._create_dir(_datadir)

        # Modules management
        self.repositories = Repositories(workdir, _datadir, self.VERSION)
        self.modules_loader = RepositoryModulesLoader(self.repositories)

        # Backend instances config
        if not backends_filename:
            backends_filename = os.environ.get('WEBOOB_BACKENDS', os.path.join(self.workdir, self.BACKENDS_FILENAME))
        elif not backends_filename.startswith('/'):
            backends_filename = os.path.join(self.workdir, backends_filename)
        self.backends_config = BackendsConfig(backends_filename)

    def _create_dir(self, name):
        if not os.path.exists(name):
            os.makedirs(name)
        elif not os.path.isdir(name):
            self.logger.error(u'"%s" is not a directory', name)

    def update(self, progress=PrintProgress()):
        """
        Update modules from repositories.
        """
        self.repositories.update(progress)

        modules_to_check = set([module_name for _, module_name, _ in self.backends_config.iter_backends()])
        for module_name in modules_to_check:
            minfo = self.repositories.get_module_info(module_name)
            if minfo and not minfo.is_installed():
                self.repositories.install(minfo, progress)

    def build_backend(self, module_name, params=None, storage=None, name=None):
        """
        Create a single backend which is not listed in configuration.

        :param module_name: name of module
        :param params: parameters to give to backend
        :type params: :class:`dict`
        :param storage: storage to use
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param name: name of backend
        :type name: :class:`basestring`
        :rtype: :class:`weboob.tools.backend.Module`
        """
        minfo = self.repositories.get_module_info(module_name)
        if minfo is None:
            raise ModuleLoadError(module_name, 'Module does not exist.')

        if not minfo.is_installed():
            self.repositories.install(minfo)

        return super(Weboob, self).build_backend(module_name, params, storage, name)

    def load_backends(self, caps=None, names=None, modules=None, exclude=None, storage=None, errors=None):
        """
        Load backends listed in config file.

        :param caps: load backends which implement all of specified caps
        :type caps: tuple[:class:`weboob.capabilities.base.Capability`]
        :param names: load backends with instance name in list
        :type names: tuple[:class:`str`]
        :param modules: load backends which module is in list
        :type modules: tuple[:class:`str`]
        :param exclude: do not load modules in list
        :type exclude: tuple[:class:`str`]
        :param storage: use this storage if specified
        :type storage: :class:`weboob.tools.storage.IStorage`
        :param errors: if specified, store every errors in this list
        :type errors: list[:class:`LoadError`]
        :returns: loaded backends
        :rtype: dict[:class:`str`, :class:`weboob.tools.backend.Module`]
        """
        loaded = {}
        if storage is None:
            storage = self.storage

        if not self.repositories.check_repositories():
            self.logger.error(u'Repositories are not consistent with the sources.list')
            raise VersionsMismatchError(u'Versions mismatch, please run "weboob-config update"')

        for instance_name, module_name, params in self.backends_config.iter_backends():
            if '_enabled' in params and not params['_enabled'].lower() in ('1', 'y', 'true', 'on', 'yes') or \
               names is not None and instance_name not in names or \
               modules is not None and module_name not in modules or \
               exclude is not None and module_name in exclude:
                continue

            minfo = self.repositories.get_module_info(module_name)
            if minfo is None:
                self.logger.warning(u'Backend "%s" is referenced in %s but was not found. '
                                    u'Perhaps a missing repository?', module_name, self.backends_config.confpath)
                continue

            if caps is not None and not minfo.has_caps(caps):
                continue

            if not minfo.is_installed():
                self.repositories.install(minfo)

            module = None
            try:
                module = self.modules_loader.get_or_load_module(module_name)
            except ModuleLoadError as e:
                self.logger.error(u'Unable to load module "%s": %s', module_name, e)
                continue

            if instance_name in self.backend_instances:
                self.logger.warning(u'Oops, the backend "%s" is already loaded. Unload it before reloading...', instance_name)
                self.unload_backends(instance_name)

            try:
                backend_instance = module.create_instance(self, instance_name, params, storage)
            except Module.ConfigError as e:
                if errors is not None:
                    errors.append(self.LoadError(instance_name, e))
            else:
                self.backend_instances[instance_name] = loaded[instance_name] = backend_instance
        return loaded
<|endoftext|>