<commit_before>from ontobio.io import assocparser, gpadparser
from ontobio import ecomap
import click
import pandas as pd
import datetime
from ontobio.io import qc
from ontobio.io.assocparser import Report
from ontobio.model import collections
import warnings
from pandas.core.common import SettingWithCopyWarning
warnings.simplefilter(action="ignore", category=SettingWithCopyWarning)

@click.command()
@click.option("--file1",
              "-file1",
              type=click.Path(),
              required=True,
              help='file1 is the source file.')
@click.option("--file2",
              "-file2",
              type=click.Path(),
              required=True,
              help='file2 is the file that is the result of a transformation.')
@click.option("--output",
              "-o",
              type=click.STRING,
              required=True,
              help='the name of the prefix for all files generated by this tool.')
@click.option("--group_by_column",
              "-gb",
              type=click.STRING,
              multiple=True,
              required=False,
              help='Options to group by include: subject, object, and/or evidence_code.'
                   'If more than one of these parameters is listed (ie: -gb = evidence_code, -gb entity_identifier, '
                   'the grouping report will group by evidence_code and entity_identifier)')
@click.option("--restrict_to_decreases",
              "-rtd",
              type=click.BOOL,
              required=False,
              help='Only report group by results when the second file shows a decrease in number by grouping column')
def compare_files(file1, file2, output, group_by_column, restrict_to_decreases):
    # decide which parser to instantiate, GAF or GPAD
    pd.set_option('display.max_rows', 35000)
    df_file1, df_file2, assocs1, assocs2 = get_parser(file1, file2)

    # get the number of counts per column of each file and summarize.
    generate_count_report(df_file1, df_file2, file1, file2, output)

    # try to figure out how many Association objects match in each file.
    compare_associations(assocs1, assocs2, output, file1, file2)

    # group_by is a list of strings exactly matching column names.
    generate_group_report(df_file1, df_file2, group_by_column, file1, file2, restrict_to_decreases, output)

def generate_count_report(df_file1, df_file2, file1, file2, output):
    file1_groups, counts_frame1 = get_column_count(df_file1, file1)
    file2_groups, counts_frame2 = get_column_count(df_file2, file2)

    merged_frame = pd.concat([counts_frame1, counts_frame2], axis=1)
    merged_frame.astype('Int64')
    merged_frame.to_csv(output + "_counts_per_column_report", sep='\t')
    s = "\n\n## COLUMN COUNT SUMMARY \n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Compared Files: " + file1 + ", " + file2 + "\n"
    s += "  * See Report File: " + output + "_counts_per_column_report" +"\n\n"
    print(s)
    print(merged_frame)

def generate_group_report(df_file1, df_file2, group_by_column, file1, file2, restrict_to_decreases, output):
    if len(group_by_column) > 0:

        s = "\n\n## GROUP BY SUMMARY \n\n"
        s += "This report generated on {}\n\n".format(datetime.date.today())
        s += "  * Group By Columns: " + str(group_by_column) + "\n"
        s += "  * Compared Files: " + file1 + ", " + file2 + "\n"

        for group in group_by_column:
            file1_groups, grouped_frame1 = get_group_by(df_file1, group, file1)
            file2_groups, grouped_frame2 = get_group_by(df_file2, group, file2)

            merged_group_frame = pd.concat([grouped_frame1, grouped_frame2], axis=1)
            merged_group_frame_no_nulls = merged_group_frame.fillna(0)
            fix_int_df = merged_group_frame_no_nulls.astype(int)
            column1 = fix_int_df.columns[0]
            column2 = fix_int_df.columns[1]+"2"
            fix_int_df.columns.values[1] = column2
            if restrict_to_decreases:
                df = fix_int_df.query("{0}".format(column1) + " > " + "{0}".format(column2))
            else:
                df = fix_int_df.query("{0}".format(column1) + " != " + "{0}".format(column2))

            s += "  * Number of unqiue " + group + "s that show differences: " + str(len(df.index)) + "\n"
            s += "  * See output file " + output + "_" + group + "_counts_per_column_report" + "\n"
            df.rename(columns={list(df)[0]: file1}, inplace=True)
            df.rename(columns={list(df)[1]: file2}, inplace=True)
            df.to_csv(output + "_" + group + "_counts_per_column_report", sep='\t')
            print(s)
            print("\n\n")

def compare_associations(assocs1, assocs2, output, file1, file2):
    compare_report_file = open(output + "_compare_report", "w")
    processed_associations = len(assocs1)

    report = Report()

    set1 = set((str(x.subject.id),
                str(x.object.id),
                normalize_relation(x.relation),
                x.negated,
                x.evidence.type,
                x.evidence._supporting_reference_to_str(),
                x.evidence._with_support_from_to_str()
                ) for x in assocs2 if type(x) != dict)
    difference = [y for y in assocs1 if type(y) != dict
                  if (str(y.subject.id),
                      str(y.object.id),
                      normalize_relation(y.relation),
                      y.negated,
                      y.evidence.type,
                      y.evidence._supporting_reference_to_str(),
                      y.evidence._with_support_from_to_str()
                      ) not in set1]

    for diff in difference:
        report.add_association(diff)
        report.n_lines = report.n_lines + 1
        report.error(diff.source_line, qc.ResultType.ERROR, "line from %s has NO match in %s" % (file1, file2), "")

    md_report, number_of_messages = markdown_report(report, processed_associations)
    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Unmatched Associations: {}\n".format(number_of_messages)
    s += "  * Total Associations Compared: " + str(len(assocs1)) + "\n"
    s += "  * See report: " + output + "_compare_report" + "\n"

    print(s)
    compare_report_file.write(md_report)
    compare_report_file.close()

def markdown_report(report, processed_lines) -> (str, str):

    json = report.to_report_json()

    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Associations Compared: " + str(processed_lines) + "\n"

    for (rule, messages) in sorted(json["messages"].items(), key=lambda t: t[0]):
        s += "### {rule}\n\n".format(rule=rule)
        s += "* total missing annotations: {amount}\n".format(amount=len(messages))
        s += "\n"
        if len(messages) > 0:
            s += "#### Messages\n\n"
        for message in messages:
            obj = " ({})".format(message["obj"]) if message["obj"] else ""
            s += "* {level} - {type}: {message}{obj} -- `{line}`\n".format(level=message["level"],
                                                                           type=message["type"],
                                                                           message=message["message"],
                                                                           line=message["line"],
                                                                           obj=obj)

        return s, len(messages)

def get_typed_parser(file_handle, filename):
    parser = assocparser.AssocParser()

    for line in file_handle:
        if assocparser.AssocParser().is_header(line):
            returned_parser = collections.create_parser_from_header(line, assocparser.AssocParserConfig())
            if returned_parser is not None:
                parser = returned_parser
        else:
            continue
    if isinstance(parser, gpadparser.GpadParser):
        df_file = read_gpad_csv(filename, parser.version)
    else:
        df_file = read_gaf_csv(filename, parser.version)

    return df_file, parser

def normalize_relation(relation: str) -> str:
    if ":" in str(relation):
        return str(relation)
    else:
        return romap.keys()[romap.values().index(str(relation))]

def get_parser(file1, file2):

    file1_obj = assocparser.AssocParser()._ensure_file(file1)
    df_file1, parser1 = get_typed_parser(file1_obj, file1)
    file2_obj = assocparser.AssocParser()._ensure_file(file2)
    df_file2, parser2 = get_typed_parser(file2_obj, file2)

    assocs1 = parser1.parse(file1)
    assocs2 = parser2.parse(file2)

    return df_file1, df_file2, assocs1, assocs2

def read_gaf_csv(filename, version):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "DB_Object_Symbol",
                                    "Qualifier",
                                    "GO_ID",
                                    "DB_Reference",
                                    "Evidence_code",
                                    "With_or_From",
                                    "Aspect",
                                    "DB_Object_Name",
                                    "DB_Object_Synonym",
                                    "DB_Object_Type,"
                                    "Taxon",
                                    "Date",
                                    "Assigned_By",
                                    "Annotation_Extension",
                                    "Gene_Product_Form_ID"]).fillna("")
    new_df = data_frame.filter(['DB_Object_ID', 'Qualifier', 'GO_ID', 'Evidence_code', 'DB_Reference'], axis=1)
    for eco_code in ecomapping.mappings():
        for ev in new_df['Evidence_code']:
            if eco_code[2] == ev:
                new_df['Evidence_code'] = new_df['Evidence_code'].replace([eco_code[2]],
                                                                              ecomapping.ecoclass_to_coderef(
                                                                                eco_code[2])[0])
    return new_df

def read_gpad_csv(filename, version):
    if version.startswith("1"):
        data_frame = pd.read_csv(filename,
                                 comment='!',
                                 sep='\t',
                                 header=None,
                                 na_filter=False,
                                 names=gpad_1_2_format).fillna("")
        df = data_frame.filter(['db', 'subject', 'qualifiers', 'relation', 'object', 'evidence_code', 'reference'], axis=1)
        concat_column = df['db'] + ":" + df['subject']
        df['concat_column'] = concat_column
        filtered_df = df.filter(['concat_column', 'qualifiers', 'relation', 'object', 'evidence_code', 'reference'])
        filtered_df.rename(columns={'concat_column': 'subject'}, inplace=True)
        new_df = filtered_df
    else:
        data_frame = pd.read_csv(filename,
                                 comment='!',
                                 sep='\t',
                                 header=None,
                                 na_filter=False,
                                 names=gpad_2_0_format).fillna("")
        new_df = data_frame.filter(['subject', 'negation', 'relation', 'object', 'evidence_code', 'reference'], axis=1)
    ecomapping = ecomap.EcoMap()
    for eco_code in ecomapping.mappings():
        for ev in new_df['evidence_code']:
            if eco_code[2] == ev:
                new_df['evidence_code'] = new_df['evidence_code'].replace([eco_code[2]],
                                                                          ecomapping.ecoclass_to_coderef(eco_code[2])[0])

    # normalize MGI ids
    config = assocparser.AssocParserConfig()
    config.remove_double_prefixes = True
    parser = gpadparser.GpadParser(config=config)
    for i, r in enumerate(new_df['subject']):
        r1 = parser._normalize_id(r)
        new_df.at[i, 'subject'] = r1
    return new_df

def get_group_by(data_frame, group, file):
    stats = {'filename': file, 'total_rows': data_frame.shape[0]}
    grouped_frame = data_frame.groupby(group)[group].count().to_frame()
    without_nulls = grouped_frame.fillna(0)
    return stats, without_nulls

def get_column_count(data_frame, file):
    stats = {'filename': file, 'total_rows': data_frame.shape[0]}
    count_frame = data_frame.nunique().to_frame(file)
    return stats, count_frame

romap = {"RO:0002327": "enables",
         "RO:0002326": "contributes_to",
         "RO:0002331": "involved_in",
         "RO:0002263": "acts_upstream_of",
         "RO:0004034": "acts_upstream_of_positive_effect",
         "RO:0004035": "acts_upstream_of_negative_effect",
         "RO:0002264": "acts_upstream_of_or_within",
         "RO:0004032": "acts_upstream_of_or_within_postitive_effect",
         "RO:0004033": "acts_upstream_of_or_within_negative_effect",
         "RO:0001025": "located_in",
         "BFO:0000050": "part_of",
         "RO:0002432": "is_active_in",
         "RO:0002325": "colocalizes_with"}

gpad_1_2_format = ["db",
                   "subject",
                   "qualifiers",
                   "object",
                   "reference",
                   "evidence_code",
                   "with_or_from",
                   "interacting_taxon",
                   "date",
                   "provided_by",
                   "annotation_extensions",
                   "properties"]

gpad_2_0_format = ["subject",
                   "negated",
                   "relation",
                   "object",
                   "reference",
                   "evidence_code",
                   "with_or_from",
                   "interacting_taxon",
                   "date",
                   "provided_by",
                   "annotation_extensions",
                   "properties"]

gaf_format = ["DB",
              "DB_Object_ID",
              "DB_Object_Symbol",
              "Qualifier",
              "GO_ID",
              "DB_Reference",
              "Evidence_code",
              "With_or_From",
              "Aspect",
              "DB_Object_Name",
              "DB_Object_Synonym",
              "DB_Object_Type",
              "Taxon",
              "Date",
              "Assigned_By",
              "Annotation_Extension",
              "Gene_Product_Form_ID"]

if __name__ == '__main__':
    compare_files()
<commit_msg>restrict group by choices to evidence_code, subject, object<commit_after>from ontobio.io import assocparser, gpadparser
from ontobio import ecomap
import click
import pandas as pd
import datetime
from ontobio.io import qc
from ontobio.io.assocparser import Report
from ontobio.model import collections
import warnings
from pandas.core.common import SettingWithCopyWarning
warnings.simplefilter(action="ignore", category=SettingWithCopyWarning)

@click.command()
@click.option("--file1",
              "-file1",
              type=click.Path(),
              required=True,
              help='file1 is the source file.')
@click.option("--file2",
              "-file2",
              type=click.Path(),
              required=True,
              help='file2 is the file that is the result of a transformation.')
@click.option("--output",
              "-o",
              type=click.STRING,
              required=True,
              help='the name of the prefix for all files generated by this tool.')
@click.option("--group_by_column",
              "-gb",
              type=click.Choice(['evidence_code', 'subject', 'object']),
              multiple=True,
              required=False,
              help='Options to group by include: subject, object, and/or evidence_code.')
@click.option("--restrict_to_decreases",
              "-rtd",
              type=click.BOOL,
              required=False,
              help='Only report group by results when the second file shows a decrease in number by grouping column')
def compare_files(file1, file2, output, group_by_column, restrict_to_decreases):
    # decide which parser to instantiate, GAF or GPAD
    pd.set_option('display.max_rows', 35000)
    df_file1, df_file2, assocs1, assocs2 = get_parser(file1, file2)

    # get the number of counts per column of each file and summarize.
    generate_count_report(df_file1, df_file2, file1, file2, output)

    # try to figure out how many Association objects match in each file.
    compare_associations(assocs1, assocs2, output, file1, file2)

    # group_by is a list of strings exactly matching column names.
    generate_group_report(df_file1, df_file2, group_by_column, file1, file2, restrict_to_decreases, output)

def generate_count_report(df_file1, df_file2, file1, file2, output):
    file1_groups, counts_frame1 = get_column_count(df_file1, file1)
    file2_groups, counts_frame2 = get_column_count(df_file2, file2)

    merged_frame = pd.concat([counts_frame1, counts_frame2], axis=1)
    merged_frame.astype('Int64')
    merged_frame.to_csv(output + "_counts_per_column_report", sep='\t')
    s = "\n\n## COLUMN COUNT SUMMARY \n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Compared Files: " + file1 + ", " + file2 + "\n"
    s += "  * See Report File: " + output + "_counts_per_column_report" +"\n\n"
    print(s)
    print(merged_frame)

def generate_group_report(df_file1, df_file2, group_by_column, file1, file2, restrict_to_decreases, output):
    if len(group_by_column) > 0:

        s = "\n\n## GROUP BY SUMMARY \n\n"
        s += "This report generated on {}\n\n".format(datetime.date.today())
        s += "  * Group By Columns: " + str(group_by_column) + "\n"
        s += "  * Compared Files: " + file1 + ", " + file2 + "\n"

        for group in group_by_column:
            file1_groups, grouped_frame1 = get_group_by(df_file1, group, file1)
            file2_groups, grouped_frame2 = get_group_by(df_file2, group, file2)

            merged_group_frame = pd.concat([grouped_frame1, grouped_frame2], axis=1)
            merged_group_frame_no_nulls = merged_group_frame.fillna(0)
            fix_int_df = merged_group_frame_no_nulls.astype(int)
            column1 = fix_int_df.columns[0]
            column2 = fix_int_df.columns[1]+"2"
            fix_int_df.columns.values[1] = column2
            if restrict_to_decreases:
                df = fix_int_df.query("{0}".format(column1) + " > " + "{0}".format(column2))
            else:
                df = fix_int_df.query("{0}".format(column1) + " != " + "{0}".format(column2))

            s += "  * Number of unqiue " + group + "s that show differences: " + str(len(df.index)) + "\n"
            s += "  * See output file " + output + "_" + group + "_counts_per_column_report" + "\n"
            df.rename(columns={list(df)[0]: file1}, inplace=True)
            df.rename(columns={list(df)[1]: file2}, inplace=True)
            df.to_csv(output + "_" + group + "_counts_per_column_report", sep='\t')
            print(s)
            print("\n\n")

def compare_associations(assocs1, assocs2, output, file1, file2):
    compare_report_file = open(output + "_compare_report", "w")
    processed_associations = len(assocs1)

    report = Report()

    set1 = set((str(x.subject.id),
                str(x.object.id),
                normalize_relation(x.relation),
                x.negated,
                x.evidence.type,
                x.evidence._supporting_reference_to_str(),
                x.evidence._with_support_from_to_str()
                ) for x in assocs2 if type(x) != dict)
    difference = [y for y in assocs1 if type(y) != dict
                  if (str(y.subject.id),
                      str(y.object.id),
                      normalize_relation(y.relation),
                      y.negated,
                      y.evidence.type,
                      y.evidence._supporting_reference_to_str(),
                      y.evidence._with_support_from_to_str()
                      ) not in set1]

    for diff in difference:
        report.add_association(diff)
        report.n_lines = report.n_lines + 1
        report.error(diff.source_line, qc.ResultType.ERROR, "line from %s has NO match in %s" % (file1, file2), "")

    md_report, number_of_messages = markdown_report(report, processed_associations)
    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Unmatched Associations: {}\n".format(number_of_messages)
    s += "  * Total Associations Compared: " + str(len(assocs1)) + "\n"
    s += "  * See report: " + output + "_compare_report" + "\n"

    print(s)
    compare_report_file.write(md_report)
    compare_report_file.close()

def markdown_report(report, processed_lines) -> (str, str):

    json = report.to_report_json()

    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Associations Compared: " + str(processed_lines) + "\n"

    for (rule, messages) in sorted(json["messages"].items(), key=lambda t: t[0]):
        s += "### {rule}\n\n".format(rule=rule)
        s += "* total missing annotations: {amount}\n".format(amount=len(messages))
        s += "\n"
        if len(messages) > 0:
            s += "#### Messages\n\n"
        for message in messages:
            obj = " ({})".format(message["obj"]) if message["obj"] else ""
            s += "* {level} - {type}: {message}{obj} -- `{line}`\n".format(level=message["level"],
                                                                           type=message["type"],
                                                                           message=message["message"],
                                                                           line=message["line"],
                                                                           obj=obj)

        return s, len(messages)

def get_typed_parser(file_handle, filename):
    parser = assocparser.AssocParser()

    for line in file_handle:
        if assocparser.AssocParser().is_header(line):
            returned_parser = collections.create_parser_from_header(line, assocparser.AssocParserConfig())
            if returned_parser is not None:
                parser = returned_parser
        else:
            continue
    if isinstance(parser, gpadparser.GpadParser):
        df_file = read_gpad_csv(filename, parser.version)
    else:
        df_file = read_gaf_csv(filename, parser.version)

    return df_file, parser

def normalize_relation(relation: str) -> str:
    if ":" in str(relation):
        return str(relation)
    else:
        return romap.keys()[romap.values().index(str(relation))]

def get_parser(file1, file2):

    file1_obj = assocparser.AssocParser()._ensure_file(file1)
    df_file1, parser1 = get_typed_parser(file1_obj, file1)
    file2_obj = assocparser.AssocParser()._ensure_file(file2)
    df_file2, parser2 = get_typed_parser(file2_obj, file2)

    assocs1 = parser1.parse(file1)
    assocs2 = parser2.parse(file2)

    return df_file1, df_file2, assocs1, assocs2

def read_gaf_csv(filename, version):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "DB_Object_Symbol",
                                    "Qualifier",
                                    "GO_ID",
                                    "DB_Reference",
                                    "Evidence_code",
                                    "With_or_From",
                                    "Aspect",
                                    "DB_Object_Name",
                                    "DB_Object_Synonym",
                                    "DB_Object_Type,"
                                    "Taxon",
                                    "Date",
                                    "Assigned_By",
                                    "Annotation_Extension",
                                    "Gene_Product_Form_ID"]).fillna("")
    new_df = data_frame.filter(['DB_Object_ID', 'Qualifier', 'GO_ID', 'Evidence_code', 'DB_Reference'], axis=1)
    for eco_code in ecomapping.mappings():
        for ev in new_df['Evidence_code']:
            if eco_code[2] == ev:
                new_df['Evidence_code'] = new_df['Evidence_code'].replace([eco_code[2]],
                                                                              ecomapping.ecoclass_to_coderef(
                                                                                eco_code[2])[0])
    return new_df

def read_gpad_csv(filename, version):
    if version.startswith("1"):
        data_frame = pd.read_csv(filename,
                                 comment='!',
                                 sep='\t',
                                 header=None,
                                 na_filter=False,
                                 names=gpad_1_2_format).fillna("")
        df = data_frame.filter(['db', 'subject', 'qualifiers', 'relation', 'object', 'evidence_code', 'reference'], axis=1)
        concat_column = df['db'] + ":" + df['subject']
        df['concat_column'] = concat_column
        filtered_df = df.filter(['concat_column', 'qualifiers', 'relation', 'object', 'evidence_code', 'reference'])
        filtered_df.rename(columns={'concat_column': 'subject'}, inplace=True)
        new_df = filtered_df
    else:
        data_frame = pd.read_csv(filename,
                                 comment='!',
                                 sep='\t',
                                 header=None,
                                 na_filter=False,
                                 names=gpad_2_0_format).fillna("")
        new_df = data_frame.filter(['subject', 'negation', 'relation', 'object', 'evidence_code', 'reference'], axis=1)
    ecomapping = ecomap.EcoMap()
    for eco_code in ecomapping.mappings():
        for ev in new_df['evidence_code']:
            if eco_code[2] == ev:
                new_df['evidence_code'] = new_df['evidence_code'].replace([eco_code[2]],
                                                                          ecomapping.ecoclass_to_coderef(eco_code[2])[0])

    # normalize MGI ids
    config = assocparser.AssocParserConfig()
    config.remove_double_prefixes = True
    parser = gpadparser.GpadParser(config=config)
    for i, r in enumerate(new_df['subject']):
        r1 = parser._normalize_id(r)
        new_df.at[i, 'subject'] = r1
    return new_df

def get_group_by(data_frame, group, file):
    stats = {'filename': file, 'total_rows': data_frame.shape[0]}
    grouped_frame = data_frame.groupby(group)[group].count().to_frame()
    without_nulls = grouped_frame.fillna(0)
    return stats, without_nulls

def get_column_count(data_frame, file):
    stats = {'filename': file, 'total_rows': data_frame.shape[0]}
    count_frame = data_frame.nunique().to_frame(file)
    return stats, count_frame

romap = {"RO:0002327": "enables",
         "RO:0002326": "contributes_to",
         "RO:0002331": "involved_in",
         "RO:0002263": "acts_upstream_of",
         "RO:0004034": "acts_upstream_of_positive_effect",
         "RO:0004035": "acts_upstream_of_negative_effect",
         "RO:0002264": "acts_upstream_of_or_within",
         "RO:0004032": "acts_upstream_of_or_within_postitive_effect",
         "RO:0004033": "acts_upstream_of_or_within_negative_effect",
         "RO:0001025": "located_in",
         "BFO:0000050": "part_of",
         "RO:0002432": "is_active_in",
         "RO:0002325": "colocalizes_with"}

gpad_1_2_format = ["db",
                   "subject",
                   "qualifiers",
                   "object",
                   "reference",
                   "evidence_code",
                   "with_or_from",
                   "interacting_taxon",
                   "date",
                   "provided_by",
                   "annotation_extensions",
                   "properties"]

gpad_2_0_format = ["subject",
                   "negated",
                   "relation",
                   "object",
                   "reference",
                   "evidence_code",
                   "with_or_from",
                   "interacting_taxon",
                   "date",
                   "provided_by",
                   "annotation_extensions",
                   "properties"]

gaf_format = ["DB",
              "DB_Object_ID",
              "DB_Object_Symbol",
              "Qualifier",
              "GO_ID",
              "DB_Reference",
              "Evidence_code",
              "With_or_From",
              "Aspect",
              "DB_Object_Name",
              "DB_Object_Synonym",
              "DB_Object_Type",
              "Taxon",
              "Date",
              "Assigned_By",
              "Annotation_Extension",
              "Gene_Product_Form_ID"]

if __name__ == '__main__':
    compare_files()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
#
#    This program 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.
#
#    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 General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

import sys
import logging

LOG_NOTSET = 'notset'
LOG_DEBUG = 'debug'
LOG_TEST = 'test'
LOG_INFO = 'info'
LOG_WARNING = 'warn'
LOG_ERROR = 'error'
LOG_CRITICAL = 'critical'

logging.TEST = logging.INFO - 5
logging.addLevelName(logging.TEST, 'TEST')

_logger = logging.getLogger(__name__)

class Logger(object):
    def __init__(self):
        _logger.warning(
            "The netsvc.Logger API shouldn't be used anymore, please "
            "use the standard `logging.getLogger` API instead.")
        super(Logger, self).__init__()

    def notifyChannel(self, name, level, msg):
        _logger.warning(
            "notifyChannel API shouldn't be used anymore, please use "
            "the standard `logging` module instead.")
        from service.web_services import common

        log = logging.getLogger(__name__ + '.deprecated.' + ustr(name))

        if level in [LOG_TEST] and not hasattr(log, level):
            fct = lambda msg, *args, **kwargs: log.log(getattr(logging, level.upper()), msg, *args, **kwargs)
            setattr(log, level, fct)

        level_method = getattr(log, level)

        if isinstance(msg, Exception):
            msg = exception_to_unicode(msg)

        try:
            msg = ustr(msg).strip()
            if level in (LOG_ERROR, LOG_CRITICAL): # and tools.config.get_misc('debug','env_info',False):
                msg = common().exp_get_server_environment() + "\n" + msg

            result = msg.split('\n')
        except UnicodeDecodeError:
            result = msg.strip().split('\n')
        try:
            if len(result)>1:
                for idx, s in enumerate(result):
                    level_method('[%02d]: %s' % (idx+1, s,))
            elif result:
                level_method(result[0])
        except IOError:
            # TODO: perhaps reset the logger streams?
            #if logrotate closes our files, we end up here..
            pass
        except Exception:
            # better ignore the exception and carry on..
            pass

    def set_loglevel(self, level, logger=None):
        if logger is not None:
            log = logging.getLogger(str(logger))
        else:
            log = logging.getLogger()
        log.setLevel(logging.INFO) # make sure next msg is printed
        log.info("Log level changed to %s" % logging.getLevelName(level))
        log.setLevel(level)

    def shutdown(self):
        logging.shutdown()

# TODO get_encodings, ustr and exception_to_unicode were originally from tools.misc.
# There are here until we refactor tools so that this module doesn't depends on tools.

def get_encodings(hint_encoding='utf-8'):
    fallbacks = {
        'latin1': 'latin9',
        'iso-8859-1': 'iso8859-15',
        'cp1252': '1252',
    }
    if hint_encoding:
        yield hint_encoding
        if hint_encoding.lower() in fallbacks:
            yield fallbacks[hint_encoding.lower()]

    # some defaults (also taking care of pure ASCII)
    for charset in ['utf8','latin1']:
        if not (hint_encoding) or (charset.lower() != hint_encoding.lower()):
            yield charset

    from locale import getpreferredencoding
    prefenc = getpreferredencoding()
    if prefenc and prefenc.lower() != 'utf-8':
        yield prefenc
        prefenc = fallbacks.get(prefenc.lower())
        if prefenc:
            yield prefenc

def ustr(value, hint_encoding='utf-8', errors='strict'):
    """This method is similar to the builtin `unicode`, except
    that it may try multiple encodings to find one that works
    for decoding `value`, and defaults to 'utf-8' first.

    :param: value: the value to convert
    :param: hint_encoding: an optional encoding that was detecte
        upstream and should be tried first to decode ``value``.
    :param str error: optional `errors` flag to pass to the unicode
        built-in to indicate how illegal character values should be
        treated: 'strict', 'ignore' or 'replace'. Passing anything
        other than 'strict' means that the first encoding tried will
        succeed, even if it's not the correct one to use, so be
        careful!
    :rtype: unicode
    :raise: UnicodeError if value cannot be coerced to unicode
    """
    if isinstance(value, Exception):
        return exception_to_unicode(value)

    if isinstance(value, unicode):
        return value

    if not isinstance(value, basestring):
        try:
            return unicode(value, errors=errors)
        except Exception:
            raise UnicodeError('unable to convert %r' % (value,))

    for ln in get_encodings(hint_encoding):
        try:
            return unicode(value, ln, errors=errors)
        except Exception:
            pass
    raise UnicodeError('unable to convert %r' % (value,))

def exception_to_unicode(e, errors='strict'):
    if (sys.version_info[:2] < (2,6)) and hasattr(e, 'message'):
        return ustr(e.message)
    if hasattr(e, 'args'):
        return "\n".join((ustr(a, errors=errors) for a in e.args))
    try:
        return unicode(e, errors=errors)
    except Exception:
        return u"Unknown message"

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>[FIX] Fix broken creation of database in previous commit - d'oh<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
#
#    This program 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.
#
#    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 General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

import sys
import logging

LOG_NOTSET = 'notset'
LOG_DEBUG = 'debug'
LOG_TEST = 'test'
LOG_INFO = 'info'
LOG_WARNING = 'warn'
LOG_ERROR = 'error'
LOG_CRITICAL = 'critical'

logging.TEST = logging.INFO - 5
logging.addLevelName(logging.TEST, 'TEST')

_logger = logging.getLogger(__name__)

class Logger(object):
    def __init__(self):
        _logger.warning(
            "The netsvc.Logger API shouldn't be used anymore, please "
            "use the standard `logging.getLogger` API instead.")
        super(Logger, self).__init__()

    def notifyChannel(self, name, level, msg):
        _logger.warning(
            "notifyChannel API shouldn't be used anymore, please use "
            "the standard `logging` module instead.")
        from service.web_services import common

        log = logging.getLogger(__name__ + '.deprecated.' + ustr(name))

        if level in [LOG_TEST] and not hasattr(log, level):
            fct = lambda msg, *args, **kwargs: log.log(getattr(logging, level.upper()), msg, *args, **kwargs)
            setattr(log, level, fct)

        level_method = getattr(log, level)

        if isinstance(msg, Exception):
            msg = exception_to_unicode(msg)

        try:
            msg = ustr(msg).strip()
            if level in (LOG_ERROR, LOG_CRITICAL): # and tools.config.get_misc('debug','env_info',False):
                msg = common().exp_get_server_environment() + "\n" + msg

            result = msg.split('\n')
        except UnicodeDecodeError:
            result = msg.strip().split('\n')
        try:
            if len(result)>1:
                for idx, s in enumerate(result):
                    level_method('[%02d]: %s' % (idx+1, s,))
            elif result:
                level_method(result[0])
        except IOError:
            # TODO: perhaps reset the logger streams?
            #if logrotate closes our files, we end up here..
            pass
        except Exception:
            # better ignore the exception and carry on..
            pass

    def set_loglevel(self, level, logger=None):
        if logger is not None:
            log = logging.getLogger(str(logger))
        else:
            log = logging.getLogger()
        log.setLevel(logging.INFO) # make sure next msg is printed
        log.info("Log level changed to %s" % logging.getLevelName(level))
        log.setLevel(level)

    def shutdown(self):
        logging.shutdown()

# TODO get_encodings, ustr and exception_to_unicode were originally from tools.misc.
# There are here until we refactor tools so that this module doesn't depends on tools.

def get_encodings(hint_encoding='utf-8'):
    fallbacks = {
        'latin1': 'latin9',
        'iso-8859-1': 'iso8859-15',
        'cp1252': '1252',
    }
    if hint_encoding:
        yield hint_encoding
        if hint_encoding.lower() in fallbacks:
            yield fallbacks[hint_encoding.lower()]

    # some defaults (also taking care of pure ASCII)
    for charset in ['utf8','latin1']:
        if not (hint_encoding) or (charset.lower() != hint_encoding.lower()):
            yield charset

    from locale import getpreferredencoding
    prefenc = getpreferredencoding()
    if prefenc and prefenc.lower() != 'utf-8':
        yield prefenc
        prefenc = fallbacks.get(prefenc.lower())
        if prefenc:
            yield prefenc

def ustr(value, hint_encoding='utf-8', errors='strict'):
    """This method is similar to the builtin `unicode`, except
    that it may try multiple encodings to find one that works
    for decoding `value`, and defaults to 'utf-8' first.

    :param: value: the value to convert
    :param: hint_encoding: an optional encoding that was detecte
        upstream and should be tried first to decode ``value``.
    :param str error: optional `errors` flag to pass to the unicode
        built-in to indicate how illegal character values should be
        treated when converting a string: 'strict', 'ignore' or 'replace'.
        Passing anything other than 'strict' means that the first
        encoding tried will be used, even if it's not the correct
        one to use, so be careful! Ignore if value is not a string/unicode.
    :rtype: unicode
    :raise: UnicodeError if value cannot be coerced to unicode
    """
    if isinstance(value, Exception):
        return exception_to_unicode(value)

    if isinstance(value, unicode):
        return value

    if not isinstance(value, basestring):
        try:
            return unicode(value)
        except Exception:
            raise UnicodeError('unable to convert %r' % (value,))

    for ln in get_encodings(hint_encoding):
        try:
            return unicode(value, ln, errors=errors)
        except Exception:
            pass
    raise UnicodeError('unable to convert %r' % (value,))

def exception_to_unicode(e):
    if (sys.version_info[:2] < (2,6)) and hasattr(e, 'message'):
        return ustr(e.message)
    if hasattr(e, 'args'):
        return "\n".join((ustr(a) for a in e.args))
    try:
        return unicode(e)
    except Exception:
        return u"Unknown message"

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2019 Matrix.org Foundation C.I.C.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import logging
import os
import sys
import tempfile

from canonicaljson import json

from twisted.internet import defer, task

import synapse
from synapse.app import _base
from synapse.config._base import ConfigError
from synapse.config.homeserver import HomeServerConfig
from synapse.config.logger import setup_logging
from synapse.handlers.admin import ExfiltrationWriter
from synapse.replication.slave.storage._base import BaseSlavedStore
from synapse.replication.slave.storage.account_data import SlavedAccountDataStore
from synapse.replication.slave.storage.appservice import SlavedApplicationServiceStore
from synapse.replication.slave.storage.client_ips import SlavedClientIpStore
from synapse.replication.slave.storage.deviceinbox import SlavedDeviceInboxStore
from synapse.replication.slave.storage.devices import SlavedDeviceStore
from synapse.replication.slave.storage.events import SlavedEventStore
from synapse.replication.slave.storage.filtering import SlavedFilteringStore
from synapse.replication.slave.storage.groups import SlavedGroupServerStore
from synapse.replication.slave.storage.presence import SlavedPresenceStore
from synapse.replication.slave.storage.push_rule import SlavedPushRuleStore
from synapse.replication.slave.storage.receipts import SlavedReceiptsStore
from synapse.replication.slave.storage.registration import SlavedRegistrationStore
from synapse.replication.slave.storage.room import RoomStore
from synapse.replication.tcp.client import ReplicationClientHandler
from synapse.server import HomeServer
from synapse.storage.engines import create_engine
from synapse.util.logcontext import LoggingContext
from synapse.util.versionstring import get_version_string

logger = logging.getLogger("synapse.app.admin_cmd")

class AdminCmdSlavedStore(
    SlavedReceiptsStore,
    SlavedAccountDataStore,
    SlavedApplicationServiceStore,
    SlavedRegistrationStore,
    SlavedFilteringStore,
    SlavedPresenceStore,
    SlavedGroupServerStore,
    SlavedDeviceInboxStore,
    SlavedDeviceStore,
    SlavedPushRuleStore,
    SlavedEventStore,
    SlavedClientIpStore,
    RoomStore,
    BaseSlavedStore,
):
    pass

class AdminCmdServer(HomeServer):
    DATASTORE_CLASS = AdminCmdSlavedStore

    def _listen_http(self, listener_config):
        pass

    def start_listening(self, listeners):
        pass

    def build_tcp_replication(self):
        return AdminCmdReplicationHandler(self)

class AdminCmdReplicationHandler(ReplicationClientHandler):
    @defer.inlineCallbacks
    def on_rdata(self, stream_name, token, rows):
        pass

    def get_streams_to_replicate(self):
        return {}

@defer.inlineCallbacks
def export_data_command(hs, args):
    """Export data for a user.

    Args:
        hs (HomeServer)
        args (argparse.Namespace)
    """

    user_id = args.user_id
    directory = args.output_directory

    res = yield hs.get_handlers().admin_handler.export_user_data(
        user_id, FileExfiltrationWriter(user_id, directory=directory)
    )
    print(res)

class FileExfiltrationWriter(ExfiltrationWriter):
    """An ExfiltrationWriter that writes the users data to a directory.
    Returns the directory location on completion.

    Note: This writes to disk on the main reactor thread.

    Args:
        user_id (str): The user whose data is being exfiltrated.
        directory (str|None): The directory to write the data to, if None then
            will write to a temporary directory.
    """

    def __init__(self, user_id, directory=None):
        self.user_id = user_id

        if directory:
            self.base_directory = directory
        else:
            self.base_directory = tempfile.mkdtemp(
                prefix="synapse-exfiltrate__%s__" % (user_id,)
            )

        os.makedirs(self.base_directory, exist_ok=True)
        if list(os.listdir(self.base_directory)):
            raise Exception("Directory must be empty")

    def write_events(self, room_id, events):
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        os.makedirs(room_directory, exist_ok=True)
        events_file = os.path.join(room_directory, "events")

        with open(events_file, "a") as f:
            for event in events:
                print(json.dumps(event.get_pdu_json()), file=f)

    def write_state(self, room_id, event_id, state):
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        state_directory = os.path.join(room_directory, "state")
        os.makedirs(state_directory, exist_ok=True)

        event_file = os.path.join(state_directory, event_id)

        with open(event_file, "a") as f:
            for event in state.values():
                print(json.dumps(event.get_pdu_json()), file=f)

    def write_invite(self, room_id, event, state):
        self.write_events(room_id, [event])

        # We write the invite state somewhere else as they aren't full events
        # and are only a subset of the state at the event.
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        os.makedirs(room_directory, exist_ok=True)

        invite_state = os.path.join(room_directory, "invite_state")

        with open(invite_state, "a") as f:
            for event in state.values():
                print(json.dumps(event), file=f)

    def finished(self):
        return self.base_directory

def start(config_options):
    parser = argparse.ArgumentParser(description="Synapse Admin Command")
    HomeServerConfig.add_arguments_to_parser(parser)

    subparser = parser.add_subparsers(
        title="Admin Commands",
        description="Choose an admin command to perform.",
        required=True,
        dest="command",
        metavar="<admin_command>",
        help="The admin command to perform.",
    )
    export_data_parser = subparser.add_parser(
        "export-data", help="Export all data for a user"
    )
    export_data_parser.add_argument("user_id", help="User to extra data from")
    export_data_parser.add_argument(
        "--output-directory",
        action="store",
        metavar="DIRECTORY",
        required=False,
        help="The directory to store the exported data in. Must be emtpy. Defaults"
        " to creating a temp directory.",
    )
    export_data_parser.set_defaults(func=export_data_command)

    try:
        config, args = HomeServerConfig.load_config_with_parser(parser, config_options)
    except ConfigError as e:
        sys.stderr.write("\n" + str(e) + "\n")
        sys.exit(1)

    if config.worker_app is not None:
        assert config.worker_app == "synapse.app.admin_cmd"

    # Update the config with some basic overrides so that don't have to specify
    # a full worker config.
    config.worker_app = "synapse.app.admin_cmd"

    if (
        not config.worker_daemonize
        and not config.worker_log_file
        and not config.worker_log_config
    ):
        # Since we're meant to be run as a "command" let's not redirect stdio
        # unless we've actually set log config.
        config.no_redirect_stdio = True

    # Explicitly disable background processes
    config.update_user_directory = False
    config.start_pushers = False
    config.send_federation = False

    setup_logging(config, use_worker_options=True)

    synapse.events.USE_FROZEN_DICTS = config.use_frozen_dicts

    database_engine = create_engine(config.database_config)

    ss = AdminCmdServer(
        config.server_name,
        db_config=config.database_config,
        config=config,
        version_string="Synapse/" + get_version_string(synapse),
        database_engine=database_engine,
    )

    ss.setup()

    # We use task.react as the basic run command as it correctly handles tearing
    # down the reactor when the deferreds resolve and setting the return value.
    # We also make sure that `_base.start` gets run before we actually run the
    # command.

    @defer.inlineCallbacks
    def run(_reactor):
        with LoggingContext("command"):
            yield _base.start(ss, [])
            yield args.func(ss, args)

    _base.start_worker_reactor(
        "synapse-admin-cmd", config, run_command=lambda: task.react(run)
    )

if __name__ == "__main__":
    with LoggingContext("main"):
        start(sys.argv[1:])
<commit_msg>Fix typo in synapse/app/admin_cmd.py<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2019 Matrix.org Foundation C.I.C.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import logging
import os
import sys
import tempfile

from canonicaljson import json

from twisted.internet import defer, task

import synapse
from synapse.app import _base
from synapse.config._base import ConfigError
from synapse.config.homeserver import HomeServerConfig
from synapse.config.logger import setup_logging
from synapse.handlers.admin import ExfiltrationWriter
from synapse.replication.slave.storage._base import BaseSlavedStore
from synapse.replication.slave.storage.account_data import SlavedAccountDataStore
from synapse.replication.slave.storage.appservice import SlavedApplicationServiceStore
from synapse.replication.slave.storage.client_ips import SlavedClientIpStore
from synapse.replication.slave.storage.deviceinbox import SlavedDeviceInboxStore
from synapse.replication.slave.storage.devices import SlavedDeviceStore
from synapse.replication.slave.storage.events import SlavedEventStore
from synapse.replication.slave.storage.filtering import SlavedFilteringStore
from synapse.replication.slave.storage.groups import SlavedGroupServerStore
from synapse.replication.slave.storage.presence import SlavedPresenceStore
from synapse.replication.slave.storage.push_rule import SlavedPushRuleStore
from synapse.replication.slave.storage.receipts import SlavedReceiptsStore
from synapse.replication.slave.storage.registration import SlavedRegistrationStore
from synapse.replication.slave.storage.room import RoomStore
from synapse.replication.tcp.client import ReplicationClientHandler
from synapse.server import HomeServer
from synapse.storage.engines import create_engine
from synapse.util.logcontext import LoggingContext
from synapse.util.versionstring import get_version_string

logger = logging.getLogger("synapse.app.admin_cmd")

class AdminCmdSlavedStore(
    SlavedReceiptsStore,
    SlavedAccountDataStore,
    SlavedApplicationServiceStore,
    SlavedRegistrationStore,
    SlavedFilteringStore,
    SlavedPresenceStore,
    SlavedGroupServerStore,
    SlavedDeviceInboxStore,
    SlavedDeviceStore,
    SlavedPushRuleStore,
    SlavedEventStore,
    SlavedClientIpStore,
    RoomStore,
    BaseSlavedStore,
):
    pass

class AdminCmdServer(HomeServer):
    DATASTORE_CLASS = AdminCmdSlavedStore

    def _listen_http(self, listener_config):
        pass

    def start_listening(self, listeners):
        pass

    def build_tcp_replication(self):
        return AdminCmdReplicationHandler(self)

class AdminCmdReplicationHandler(ReplicationClientHandler):
    @defer.inlineCallbacks
    def on_rdata(self, stream_name, token, rows):
        pass

    def get_streams_to_replicate(self):
        return {}

@defer.inlineCallbacks
def export_data_command(hs, args):
    """Export data for a user.

    Args:
        hs (HomeServer)
        args (argparse.Namespace)
    """

    user_id = args.user_id
    directory = args.output_directory

    res = yield hs.get_handlers().admin_handler.export_user_data(
        user_id, FileExfiltrationWriter(user_id, directory=directory)
    )
    print(res)

class FileExfiltrationWriter(ExfiltrationWriter):
    """An ExfiltrationWriter that writes the users data to a directory.
    Returns the directory location on completion.

    Note: This writes to disk on the main reactor thread.

    Args:
        user_id (str): The user whose data is being exfiltrated.
        directory (str|None): The directory to write the data to, if None then
            will write to a temporary directory.
    """

    def __init__(self, user_id, directory=None):
        self.user_id = user_id

        if directory:
            self.base_directory = directory
        else:
            self.base_directory = tempfile.mkdtemp(
                prefix="synapse-exfiltrate__%s__" % (user_id,)
            )

        os.makedirs(self.base_directory, exist_ok=True)
        if list(os.listdir(self.base_directory)):
            raise Exception("Directory must be empty")

    def write_events(self, room_id, events):
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        os.makedirs(room_directory, exist_ok=True)
        events_file = os.path.join(room_directory, "events")

        with open(events_file, "a") as f:
            for event in events:
                print(json.dumps(event.get_pdu_json()), file=f)

    def write_state(self, room_id, event_id, state):
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        state_directory = os.path.join(room_directory, "state")
        os.makedirs(state_directory, exist_ok=True)

        event_file = os.path.join(state_directory, event_id)

        with open(event_file, "a") as f:
            for event in state.values():
                print(json.dumps(event.get_pdu_json()), file=f)

    def write_invite(self, room_id, event, state):
        self.write_events(room_id, [event])

        # We write the invite state somewhere else as they aren't full events
        # and are only a subset of the state at the event.
        room_directory = os.path.join(self.base_directory, "rooms", room_id)
        os.makedirs(room_directory, exist_ok=True)

        invite_state = os.path.join(room_directory, "invite_state")

        with open(invite_state, "a") as f:
            for event in state.values():
                print(json.dumps(event), file=f)

    def finished(self):
        return self.base_directory

def start(config_options):
    parser = argparse.ArgumentParser(description="Synapse Admin Command")
    HomeServerConfig.add_arguments_to_parser(parser)

    subparser = parser.add_subparsers(
        title="Admin Commands",
        description="Choose an admin command to perform.",
        required=True,
        dest="command",
        metavar="<admin_command>",
        help="The admin command to perform.",
    )
    export_data_parser = subparser.add_parser(
        "export-data", help="Export all data for a user"
    )
    export_data_parser.add_argument("user_id", help="User to extra data from")
    export_data_parser.add_argument(
        "--output-directory",
        action="store",
        metavar="DIRECTORY",
        required=False,
        help="The directory to store the exported data in. Must be empty. Defaults"
        " to creating a temp directory.",
    )
    export_data_parser.set_defaults(func=export_data_command)

    try:
        config, args = HomeServerConfig.load_config_with_parser(parser, config_options)
    except ConfigError as e:
        sys.stderr.write("\n" + str(e) + "\n")
        sys.exit(1)

    if config.worker_app is not None:
        assert config.worker_app == "synapse.app.admin_cmd"

    # Update the config with some basic overrides so that don't have to specify
    # a full worker config.
    config.worker_app = "synapse.app.admin_cmd"

    if (
        not config.worker_daemonize
        and not config.worker_log_file
        and not config.worker_log_config
    ):
        # Since we're meant to be run as a "command" let's not redirect stdio
        # unless we've actually set log config.
        config.no_redirect_stdio = True

    # Explicitly disable background processes
    config.update_user_directory = False
    config.start_pushers = False
    config.send_federation = False

    setup_logging(config, use_worker_options=True)

    synapse.events.USE_FROZEN_DICTS = config.use_frozen_dicts

    database_engine = create_engine(config.database_config)

    ss = AdminCmdServer(
        config.server_name,
        db_config=config.database_config,
        config=config,
        version_string="Synapse/" + get_version_string(synapse),
        database_engine=database_engine,
    )

    ss.setup()

    # We use task.react as the basic run command as it correctly handles tearing
    # down the reactor when the deferreds resolve and setting the return value.
    # We also make sure that `_base.start` gets run before we actually run the
    # command.

    @defer.inlineCallbacks
    def run(_reactor):
        with LoggingContext("command"):
            yield _base.start(ss, [])
            yield args.func(ss, args)

    _base.start_worker_reactor(
        "synapse-admin-cmd", config, run_command=lambda: task.react(run)
    )

if __name__ == "__main__":
    with LoggingContext("main"):
        start(sys.argv[1:])
<|endoftext|>


<commit_before>from __future__ import absolute_import
import datetime
import json
import logging
import os
import pprint
import subprocess
import time
import tempfile
import textwrap
import humanize
import logic.symbolic
import natsort
import networkx as nx
from openpromela import bitvector
import psutil

logger = logging.getLogger(__name__)
slugs_log = logging.getLogger(__name__ + '.slugs')
SLUGS_SPEC = 'spec_slugs.txt'
BDD_FILE = 'bdd.txt'

def synthesize(spec, symbolic=True, bddfile=None, real=True):
    """Return strategy satisfying the specification `spec`.

    @param spec: GR(1) specification
    @type spec: `tulip.spec.GRSpec`

    @return: If realizable return synthesized strategy, otherwise `None`.
    @rtype: `networkx.DiGraph`
    """
    logger.info('++ compile LTL to slugsin\n')
    aut = logic.symbolic.bitblast_spec(spec)
    s = _to_slugs(aut)
    # dump for use in manual debugging
    if logger.getEffectiveLevel() < logging.DEBUG:
        with open(SLUGS_SPEC, 'w') as f:
            w = textwrap.fill(s, replace_whitespace=False)
            f.write(w)
    # call slugs
    with tempfile.NamedTemporaryFile(delete=False) as fin:
        fin.write(s)
    logger.info('\n\n spec:\n\n {spec}'.format(
        spec=spec) + '\n\n slugs in:\n\n {s}\n'.format(s=s))
    logger.info('-- done compiling to slugsin')
    realizable, out = _call_slugs(
        fin.name, symbolic=symbolic, bddfile=bddfile, real=real)
    os.unlink(fin.name)
    logger.debug('slugs output:\n{out}'.format(out=out))
    if realizable:
        logger.info('realizable')
    else:
        logger.info('not realizable')
    # symbolic strategies not loaded from file yet
    if symbolic:
        return realizable
    if not realizable:
        return None
    g = load_strategy(out)
    logger.debug(
        ('loaded strategy with vertices:\n  {v}\n'
         'and edges:\n {e}\n').format(
            v='\n  '.join(str(x) for x in g.nodes(data=True)),
            e=g.edges()))
    h = bitvector.bitfield_to_int_states(g, aut.vars)
    mealy = strategy_to_mealy(h, spec)
    return mealy

def load_strategy(filename):
    """Return `networkx.DiGraph` for strategy in JSON file."""
    dout = json.loads(filename)
    # use nx graph to represent strategy
    g = nx.DiGraph()
    dvars = dout['variables']
    for stru, d in dout['nodes'].iteritems():
        u = int(stru)
        state = dict(zip(dvars, d['state']))
        g.add_node(u, state=state)
        for v in d['trans']:
            g.add_edge(u, v)
    return g

def _to_slugs(aut):
    """Return spec in `slugsin` format.

    @type aut: `symbolic.Automaton`.
    """
    dbits = bitvector.list_bits(aut.vars)
    logger.debug(
        'slugs variables:\n{v}'.format(v=pprint.pformat(dbits)))
    f = _slugs_str
    return (
        _format_slugs_vars(dbits, 'env', 'INPUT') +
        _format_slugs_vars(dbits, 'sys', 'OUTPUT') +
        # env
        f(aut.init['env'], 'ENV_INIT') +
        f(aut.action['env'], 'ENV_TRANS') +
        f(aut.win['env'], 'ENV_LIVENESS') +
        # sys
        f(aut.init['sys'], 'SYS_INIT') +
        f(aut.action['sys'], 'SYS_TRANS') +
        f(aut.win['sys'], 'SYS_LIVENESS'))

def _slugs_str(r, name, sep='\n'):
    if not r:
        return '[{name}]\n'.format(name=name)
    sep = ' {sep} '.format(sep=sep)
    f = sep.join(x for x in r if x)
    return '[{name}]\n{f}\n\n'.format(name=name, f=f)

def _format_slugs_vars(dvars, owner, name):
    a = [var for var, d in dvars.iteritems()
         if d['owner'] == owner]
    a = natsort.natsorted(a)
    return '[{name}]\n{vars}\n\n'.format(name=name, vars='\n'.join(a))

def _call_slugs(filename, symbolic=True, bddfile=None, real=True):
    """Call `slugs` and log memory usage and time."""
    if bddfile is None:
        bddfile = BDD_FILE
    options = ['slugs', filename]
    if real:
        if symbolic:
            options.extend(['--symbolicStrategy', bddfile])
        else:
            options.append('--jsonOutput')
    else:
        options.append('--onlyRealizability')
    logger.debug('Calling: {cmd}'.format(cmd=' '.join(options)))
    try:
        p = subprocess.Popen(
            options,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    except OSError as e:
        if e.errno == os.errno.ENOENT:
            raise Exception('slugs not found in path.')
        else:
            raise
    proc = psutil.Process(p.pid)
    slugs_log.info('call slugs')
    while p.poll() is None:
        try:
            user, system = proc.cpu_times()
            rss, vms = proc.memory_info()
            t = user + system
            dt = datetime.timedelta(seconds=t)
            s = 'time: {dt}, rss: {rss}, vms: {vms}'.format(
                dt=dt,
                rss=humanize.naturalsize(rss),
                vms=humanize.naturalsize(vms))
            slugs_log.info(s)
        except psutil.AccessDenied:
            logger.debug('slugs has terminated already.')
        time.sleep(1.0)
    slugs_log.info('slugs returned')
    out, err = p.communicate()
    msg = (
        '\n slugs return code: {c}\n\n'.format(c=p.returncode) +
        '\n slugs stdrr: {c}\n\n'.format(c=err) +
        '\n slugs stdout:\n\n {out}\n\n'.format(out=out))
    logger.debug(msg)
    # error ?
    if p.returncode != 0:
        raise Exception(msg)
    realizable = 'Specification is realizable' in err
    # check sanity
    if not realizable:
        assert 'Specification is unrealizable' in err
    return realizable, out
<commit_msg>update docstring<commit_after>"""Interface to `slugs` synthesizer."""
from __future__ import absolute_import
import datetime
import json
import logging
import os
import pprint
import subprocess
import time
import tempfile
import textwrap
import humanize
import logic.symbolic
import natsort
import networkx as nx
from openpromela import bitvector
import psutil

logger = logging.getLogger(__name__)
slugs_log = logging.getLogger(__name__ + '.slugs')
SLUGS_SPEC = 'spec_slugs.txt'
BDD_FILE = 'bdd.txt'

def synthesize(spec, symbolic=True, bddfile=None, real=True):
    """Return strategy satisfying the specification `spec`.

    @param spec: specification
    @type spec: `symbolic.Automaton`

    @return: If realizable return synthesized strategy, otherwise `None`.
    @rtype: `automata.Transducer` or `symbolic.Automaton`
    """
    logger.info('++ compile LTL to slugsin\n')
    aut = logic.symbolic.bitblast_spec(spec)
    s = _to_slugs(aut)
    # dump for use in manual debugging
    if logger.getEffectiveLevel() < logging.DEBUG:
        with open(SLUGS_SPEC, 'w') as f:
            w = textwrap.fill(s, replace_whitespace=False)
            f.write(w)
    # call slugs
    with tempfile.NamedTemporaryFile(delete=False) as fin:
        fin.write(s)
    logger.info('\n\n spec:\n\n {spec}'.format(
        spec=spec) + '\n\n slugs in:\n\n {s}\n'.format(s=s))
    logger.info('-- done compiling to slugsin')
    realizable, out = _call_slugs(
        fin.name, symbolic=symbolic, bddfile=bddfile, real=real)
    os.unlink(fin.name)
    logger.debug('slugs output:\n{out}'.format(out=out))
    if realizable:
        logger.info('realizable')
    else:
        logger.info('not realizable')
    # symbolic strategies not loaded from file yet
    if symbolic:
        return realizable
    if not realizable:
        return None
    g = load_strategy(out)
    logger.debug(
        ('loaded strategy with vertices:\n  {v}\n'
         'and edges:\n {e}\n').format(
            v='\n  '.join(str(x) for x in g.nodes(data=True)),
            e=g.edges()))
    h = bitvector.bitfield_to_int_states(g, aut.vars)
    mealy = strategy_to_mealy(h, spec)
    return mealy

def load_strategy(filename):
    """Return `networkx.DiGraph` for strategy in JSON file."""
    dout = json.loads(filename)
    # use nx graph to represent strategy
    g = nx.DiGraph()
    dvars = dout['variables']
    for stru, d in dout['nodes'].iteritems():
        u = int(stru)
        state = dict(zip(dvars, d['state']))
        g.add_node(u, state=state)
        for v in d['trans']:
            g.add_edge(u, v)
    return g

def _to_slugs(aut):
    """Return spec in `slugsin` format.

    @type aut: `symbolic.Automaton`.
    """
    dbits = bitvector.list_bits(aut.vars)
    logger.debug(
        'slugs variables:\n{v}'.format(v=pprint.pformat(dbits)))
    f = _slugs_str
    return (
        _format_slugs_vars(dbits, 'env', 'INPUT') +
        _format_slugs_vars(dbits, 'sys', 'OUTPUT') +
        # env
        f(aut.init['env'], 'ENV_INIT') +
        f(aut.action['env'], 'ENV_TRANS') +
        f(aut.win['env'], 'ENV_LIVENESS') +
        # sys
        f(aut.init['sys'], 'SYS_INIT') +
        f(aut.action['sys'], 'SYS_TRANS') +
        f(aut.win['sys'], 'SYS_LIVENESS'))

def _slugs_str(r, name, sep='\n'):
    if not r:
        return '[{name}]\n'.format(name=name)
    sep = ' {sep} '.format(sep=sep)
    f = sep.join(x for x in r if x)
    return '[{name}]\n{f}\n\n'.format(name=name, f=f)

def _format_slugs_vars(dvars, owner, name):
    a = [var for var, d in dvars.iteritems()
         if d['owner'] == owner]
    a = natsort.natsorted(a)
    return '[{name}]\n{vars}\n\n'.format(name=name, vars='\n'.join(a))

def _call_slugs(filename, symbolic=True, bddfile=None, real=True):
    """Call `slugs` and log memory usage and time."""
    if bddfile is None:
        bddfile = BDD_FILE
    options = ['slugs', filename]
    if real:
        if symbolic:
            options.extend(['--symbolicStrategy', bddfile])
        else:
            options.append('--jsonOutput')
    else:
        options.append('--onlyRealizability')
    logger.debug('Calling: {cmd}'.format(cmd=' '.join(options)))
    try:
        p = subprocess.Popen(
            options,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    except OSError as e:
        if e.errno == os.errno.ENOENT:
            raise Exception('slugs not found in path.')
        else:
            raise
    proc = psutil.Process(p.pid)
    slugs_log.info('call slugs')
    while p.poll() is None:
        try:
            user, system = proc.cpu_times()
            rss, vms = proc.memory_info()
            t = user + system
            dt = datetime.timedelta(seconds=t)
            s = 'time: {dt}, rss: {rss}, vms: {vms}'.format(
                dt=dt,
                rss=humanize.naturalsize(rss),
                vms=humanize.naturalsize(vms))
            slugs_log.info(s)
        except psutil.AccessDenied:
            logger.debug('slugs has terminated already.')
        time.sleep(1.0)
    slugs_log.info('slugs returned')
    out, err = p.communicate()
    msg = (
        '\n slugs return code: {c}\n\n'.format(c=p.returncode) +
        '\n slugs stdrr: {c}\n\n'.format(c=err) +
        '\n slugs stdout:\n\n {out}\n\n'.format(out=out))
    logger.debug(msg)
    # error ?
    if p.returncode != 0:
        raise Exception(msg)
    realizable = 'Specification is realizable' in err
    # check sanity
    if not realizable:
        assert 'Specification is unrealizable' in err
    return realizable, out
<|endoftext|>


<commit_before>"""
The ``masks`` module provides filters for working with binary images.

===============================================================================
Overview
===============================================================================
Functions provided include binary dilation and erosion. These allow for N-D
array processing.

===============================================================================
API
===============================================================================
"""

__author__ = "John Kirkham <kirkhamj@janelia.hhmi.org>"
__date__ = "$Mar 31, 2015 20:47:37 EDT$"

import numpy

import scipy
import scipy.ndimage
import scipy.ndimage.filters

# Need in order to have logging information no matter what.
from nanshe.util import prof

# Get the logger
trace_logger = prof.getTraceLogger(__name__)

@prof.log_call(trace_logger)
def binary_dilation(input_array, footprint, out=None):
    """
        Performs simple binary dilation on a bool array of arbitrary dimension.

        Args:
            input_array(numpy.ndarray):         the bool array to perform dilation on.
            footprint(numpy.ndarray):           the footprint to use for the kernel.
            out(numpy.ndarray):                 a place to store the result if provided. (optional)

        Returns:
            out(numpy.ndarray):                 Same as out if out was provided.

        >>> a = numpy.array([[ True,  True, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False,  True, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False]], dtype=bool)
        >>> b = numpy.zeros_like(a)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool))
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool), out=b)
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> b
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool), out=a)
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> a
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
    """

    assert issubclass(input_array.dtype.type, (bool, numpy.bool_))

    if out is None:
        out = input_array.copy()
    elif id(input_array) != id(out):
        assert issubclass(out.dtype.type, (bool, numpy.bool_))

    scipy.ndimage.filters.maximum_filter(input_array, footprint=footprint, output=out)

    return(out)

@prof.log_call(trace_logger)
def binary_erosion(input_array, footprint, out=None):
    """
        Performs simple binary erosion on a bool array of arbitrary dimension.

        Args:
            input_array(numpy.ndarray):         the bool array to perform erosion on.
            footprint(numpy.ndarray):           the footprint to use for the kernel.
            out(numpy.ndarray):                 a place to store the result if provided. (optional)

        Returns:
            out(numpy.ndarray):                 Same as out if out was provided.

        >>> a = numpy.array([[ True,  True,  True, False, False, False,  True],
        ...                  [ True,  True,  True, False, False,  True,  True],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False,  True,  True,  True, False],
        ...                  [False,  True, False,  True,  True,  True, False],
        ...                  [False,  True, False,  True,  True,  True, False],
        ...                  [False, False, False, False, False, False,  True]], dtype=bool)
        >>> b = numpy.zeros_like(a)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool))
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool), out=b)
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> b
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool), out=a)
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> a
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
    """

    assert issubclass(input_array.dtype.type, (bool, numpy.bool_))

    if out is None:
        out = input_array.copy()
    elif id(input_array) != id(out):
        assert issubclass(out.dtype.type, (bool, numpy.bool_))

    scipy.ndimage.filters.minimum_filter(input_array, footprint=footprint, output=out)

    return(out)
<commit_msg>nanshe/imp/filters/masks.py: Wrap lines.<commit_after>"""
The ``masks`` module provides filters for working with binary images.

===============================================================================
Overview
===============================================================================
Functions provided include binary dilation and erosion. These allow for N-D
array processing.

===============================================================================
API
===============================================================================
"""

__author__ = "John Kirkham <kirkhamj@janelia.hhmi.org>"
__date__ = "$Mar 31, 2015 20:47:37 EDT$"

import numpy

import scipy
import scipy.ndimage
import scipy.ndimage.filters

# Need in order to have logging information no matter what.
from nanshe.util import prof

# Get the logger
trace_logger = prof.getTraceLogger(__name__)

@prof.log_call(trace_logger)
def binary_dilation(input_array, footprint, out=None):
    """
        Performs simple binary dilation on a bool array of arbitrary dimension.

        Args:
            input_array(numpy.ndarray):         the bool array to perform dilation on.
            footprint(numpy.ndarray):           the footprint to use for the kernel.
            out(numpy.ndarray):                 a place to store the result if provided. (optional)

        Returns:
            out(numpy.ndarray):                 Same as out if out was provided.

        >>> a = numpy.array([[ True,  True, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False,  True, False, False],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False, False, False, False, False]], dtype=bool)
        >>> b = numpy.zeros_like(a)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool))
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool), out=b)
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> b
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_dilation(a, numpy.ones(a.ndim*(3,), dtype=bool), out=a)
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> a
        array([[ True,  True,  True, False, False, False, False],
               [ True,  True,  True, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False,  True,  True,  True, False],
               [False, False, False, False, False, False, False]], dtype=bool)
    """

    assert issubclass(input_array.dtype.type, (bool, numpy.bool_))

    if out is None:
        out = input_array.copy()
    elif id(input_array) != id(out):
        assert issubclass(out.dtype.type, (bool, numpy.bool_))

    scipy.ndimage.filters.maximum_filter(
        input_array, footprint=footprint, output=out
    )

    return(out)

@prof.log_call(trace_logger)
def binary_erosion(input_array, footprint, out=None):
    """
        Performs simple binary erosion on a bool array of arbitrary dimension.

        Args:
            input_array(numpy.ndarray):         the bool array to perform erosion on.
            footprint(numpy.ndarray):           the footprint to use for the kernel.
            out(numpy.ndarray):                 a place to store the result if provided. (optional)

        Returns:
            out(numpy.ndarray):                 Same as out if out was provided.

        >>> a = numpy.array([[ True,  True,  True, False, False, False,  True],
        ...                  [ True,  True,  True, False, False,  True,  True],
        ...                  [False, False, False, False, False, False, False],
        ...                  [False, False, False,  True,  True,  True, False],
        ...                  [False,  True, False,  True,  True,  True, False],
        ...                  [False,  True, False,  True,  True,  True, False],
        ...                  [False, False, False, False, False, False,  True]], dtype=bool)
        >>> b = numpy.zeros_like(a)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool))
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool), out=b)
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> b
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)

        >>> binary_erosion(a, numpy.ones(a.ndim*(3,), dtype=bool), out=a)
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
        >>> a
        array([[ True,  True, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False,  True, False, False],
               [False, False, False, False, False, False, False],
               [False, False, False, False, False, False, False]], dtype=bool)
    """

    assert issubclass(input_array.dtype.type, (bool, numpy.bool_))

    if out is None:
        out = input_array.copy()
    elif id(input_array) != id(out):
        assert issubclass(out.dtype.type, (bool, numpy.bool_))

    scipy.ndimage.filters.minimum_filter(
        input_array, footprint=footprint, output=out
    )

    return(out)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2014 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ._base import SQLBaseStore

class StateStore(SQLBaseStore):
    """ Keeps track of the state at a given event.

    This is done by the concept of `state groups`. Every event is a assigned
    a state group (identified by an arbitrary string), which references a
    collection of state events. The current state of an event is then the
    collection of state events referenced by the event's state group.

    Hence, every change in the current state causes a new state group to be
    generated. However, if no change happens (e.g., if we get a message event
    with only one parent it inherits the state group from its parent.)

    There are three tables:
      * `state_groups`: Stores group name, first event with in the group and
        room id.
      * `event_to_state_groups`: Maps events to state groups.
      * `state_groups_state`: Maps state group to state events.
    """

    def get_state_groups(self, event_ids):
        """ Get the state groups for the given list of event_ids

        The return value is a dict mapping group names to lists of events.
        """

        def f(txn):
            groups = set()
            for event_id in event_ids:
                group = self._simple_select_one_onecol_txn(
                    txn,
                    table="event_to_state_groups",
                    keyvalues={"event_id": event_id},
                    retcol="state_group",
                    allow_none=True,
                )
                if group:
                    groups.add(group)

            res = {}
            for group in groups:
                state_ids = self._simple_select_onecol_txn(
                    txn,
                    table="state_groups_state",
                    keyvalues={"state_group": group},
                    retcol="event_id",
                )
                state = self._get_events_txn(txn, state_ids)

                res[group] = state

            return res

        return self.runInteraction(
            "get_state_groups",
            f,
        )

    def store_state_groups(self, event):
        return self.runInteraction(
            "store_state_groups",
            self._store_state_groups_txn, event
        )

    def _store_state_groups_txn(self, txn, event, context):
        if context.current_state is None:
            return

        state_events = context.current_state

        if event.is_state():
            state_events[(event.type, event.state_key)] = event

        state_group = context.state_group
        if not state_group:
            state_group = self._simple_insert_txn(
                txn,
                table="state_groups",
                values={
                    "room_id": event.room_id,
                    "event_id": event.event_id,
                },
                or_ignore=True,
            )

            for state in state_events.values():
                self._simple_insert_txn(
                    txn,
                    table="state_groups_state",
                    values={
                        "state_group": state_group,
                        "room_id": state.room_id,
                        "type": state.type,
                        "state_key": state.state_key,
                        "event_id": state.event_id,
                    },
                    or_ignore=True,
                )

        self._simple_insert_txn(
            txn,
            table="event_to_state_groups",
            values={
                "state_group": state_group,
                "event_id": event.event_id,
            },
            or_replace=True,
        )
<commit_msg>Add some debug logging<commit_after># -*- coding: utf-8 -*-
# Copyright 2014 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ._base import SQLBaseStore

import logging

logger = logging.getLogger(__name__)

class StateStore(SQLBaseStore):
    """ Keeps track of the state at a given event.

    This is done by the concept of `state groups`. Every event is a assigned
    a state group (identified by an arbitrary string), which references a
    collection of state events. The current state of an event is then the
    collection of state events referenced by the event's state group.

    Hence, every change in the current state causes a new state group to be
    generated. However, if no change happens (e.g., if we get a message event
    with only one parent it inherits the state group from its parent.)

    There are three tables:
      * `state_groups`: Stores group name, first event with in the group and
        room id.
      * `event_to_state_groups`: Maps events to state groups.
      * `state_groups_state`: Maps state group to state events.
    """

    def get_state_groups(self, event_ids):
        """ Get the state groups for the given list of event_ids

        The return value is a dict mapping group names to lists of events.
        """

        def f(txn):
            groups = set()
            for event_id in event_ids:
                group = self._simple_select_one_onecol_txn(
                    txn,
                    table="event_to_state_groups",
                    keyvalues={"event_id": event_id},
                    retcol="state_group",
                    allow_none=True,
                )
                if group:
                    groups.add(group)

            logger.debug("Got groups: %s", groups)

            res = {}
            for group in groups:
                state_ids = self._simple_select_onecol_txn(
                    txn,
                    table="state_groups_state",
                    keyvalues={"state_group": group},
                    retcol="event_id",
                )

                logger.debug(
                    "Got %d events for group %s",
                    len(state_ids), group
                )

                state = self._get_events_txn(txn, state_ids)

                res[group] = state

            return res

        return self.runInteraction(
            "get_state_groups",
            f,
        )

    def store_state_groups(self, event):
        return self.runInteraction(
            "store_state_groups",
            self._store_state_groups_txn, event
        )

    def _store_state_groups_txn(self, txn, event, context):
        if context.current_state is None:
            return

        state_events = context.current_state

        if event.is_state():
            state_events[(event.type, event.state_key)] = event

        state_group = context.state_group
        if not state_group:
            state_group = self._simple_insert_txn(
                txn,
                table="state_groups",
                values={
                    "room_id": event.room_id,
                    "event_id": event.event_id,
                },
                or_ignore=True,
            )

            for state in state_events.values():
                self._simple_insert_txn(
                    txn,
                    table="state_groups_state",
                    values={
                        "state_group": state_group,
                        "room_id": state.room_id,
                        "type": state.type,
                        "state_key": state.state_key,
                        "event_id": state.event_id,
                    },
                    or_ignore=True,
                )

        self._simple_insert_txn(
            txn,
            table="event_to_state_groups",
            values={
                "state_group": state_group,
                "event_id": event.event_id,
            },
            or_replace=True,
        )
<|endoftext|>