<commit_before>import argparse
import sys

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, aliased, joinedload
from sqlalchemy.sql.expression import true, false

from irco import graphs, models


def write(fh, graph, format=None):
    from networkx.readwrite import gexf
    gexf.write_gexf(graph, fh)


def main():
    graph_choices = graphs.get_graphs()
    argparser = argparse.ArgumentParser('irco-graph')
    argparser.add_argument('-v', '--verbose', action='store_true')
    argparser.add_argument('-y', '--years')
    argparser.add_argument('graph_type', choices=graph_choices)
    argparser.add_argument('database')
    argparser.add_argument('output', default='-', nargs='?')

    args = argparser.parse_args()
    graph_factory = graphs.get_graph(args.graph_type)

    engine = create_engine(args.database, echo=args.verbose)
    Session = sessionmaker(bind=engine)
    session = Session()

    publications = (session
                    .query(models.Publication)
                    .options(
                        joinedload('affiliations').joinedload('author'),
                        joinedload('affiliations').joinedload('institution'),
                    ))

    if args.years:
        criteria = false()
        for y in args.years.split(','):
            if '-' in y:
                start, end = y.split('-')
                if start:
                    c = models.Publication.year >= int(start)
                else:
                    c = true()

                if end:
                    c = c & (models.Publication.year <= int(end))
            else:
                c = models.Publication.year == int(y)
            criteria = criteria | c

        publications = publications.filter(criteria)


    graph = graph_factory.create(session, publications)

    if args.output == '-':
        write(sys.stdout, graph)
    else:
        with open(args.output, 'wb') as fh:
            write(fh, graph)
<commit_msg>Implement filtering by corresponding author country.<commit_after>import argparse
import sys

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, aliased, joinedload
from sqlalchemy.sql.expression import true, false

from irco import graphs, models


def write(fh, graph, format=None):
    from networkx.readwrite import gexf
    gexf.write_gexf(graph, fh)


def main():
    graph_choices = graphs.get_graphs()
    argparser = argparse.ArgumentParser('irco-graph')
    argparser.add_argument('-v', '--verbose', action='store_true')
    argparser.add_argument('-y', '--years')
    argparser.add_argument('-c', '--ca-country', action='append',
                           dest='ca_countries')
    argparser.add_argument('graph_type', choices=graph_choices)
    argparser.add_argument('database')
    argparser.add_argument('output', default='-', nargs='?')

    args = argparser.parse_args()
    graph_factory = graphs.get_graph(args.graph_type)

    engine = create_engine(args.database, echo=args.verbose)
    Session = sessionmaker(bind=engine)
    session = Session()

    publications = (session
                    .query(models.Publication)
                    .options(
                        joinedload('affiliations').joinedload('author'),
                        joinedload('affiliations').joinedload('institution'),
                    ))

    if args.years:
        criteria = false()
        for y in args.years.split(','):
            if '-' in y:
                start, end = y.split('-')
                if start:
                    c = models.Publication.year >= int(start)
                else:
                    c = true()

                if end:
                    c = c & (models.Publication.year <= int(end))
            else:
                c = models.Publication.year == int(y)
            criteria = criteria | c

        publications = publications.filter(criteria)

    if args.ca_countries:
        c_author = aliased(models.AffiliatedAuthor)
        c_institution = aliased(models.Institution)

        criteria = false()

        for country in args.ca_countries:
            criteria = criteria | (c_institution.country == country)

        publications = (publications
                        .join(c_author)
                        .join(c_institution)
                        .filter(c_author.is_corresponding == True)
                        .filter(criteria))

    graph = graph_factory.create(session, publications)

    if args.output == '-':
        write(sys.stdout, graph)
    else:
        with open(args.output, 'wb') as fh:
            write(fh, graph)

{'fraction_non_alphanumeric': 0.11326086956521739, 'fraction_numerical': 0.0, 'mean_word_length': 2.0983164983164984, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2061125', 'n_tokens_mistral': 1324, 'n_tokens_neox': 1262, 'n_words': 272}