nopaque/app/corpora/events.py

from app import db, logger, socketio
from app.events import connected_sessions
from app.models import Corpus
from flask import current_app, request
from flask_login import current_user, login_required
from .CQiWrapper.CQiClient import CQiClient
from .CQiWrapper.CQi import CONST_FIELD_MATCH, CONST_FIELD_MATCHEND


'''
' A dictionary containing lists of, with corpus ids associated, Socket.IO
' session ids (sid). {<corpus_id>: [<sid>, ...], ...}
'''
analysis_sessions = {}
'''
' A dictionary containing Socket.IO session id - CQi client pairs.
' {<sid>: CQiClient, ...}
'''
analysis_clients = {}


@socketio.on('init_corpus_analysis')
@login_required
def init_corpus_analysis(corpus_id):
    corpus = Corpus.query.get(corpus_id)
    if corpus is None:
        socketio.emit('init_corpus_analysis', '[ERROR 404]: Not Found',
                      room=request.sid)
    if not (corpus.creator == current_user or current_user.is_administrator()):
        socketio.emit('init_corpus_analysis', '[ERROR 403]: Forbidden',
                      room=request.sid)
    if corpus_id not in analysis_sessions:
        analysis_sessions[corpus_id] = [request.sid]
    else:
        analysis_sessions[corpus_id].append(request.sid)
    socketio.start_background_task(observe_corpus_analysis_connection,
                                   current_app._get_current_object(),
                                   corpus_id, request.sid)


@socketio.on('query')
@login_required
def recv_query(message):
    logger.warning(message)
    analysis_client = analysis_clients[request.sid]
    """ Prepare and execute a query """
    corpus = 'CORPUS'
    query = message['query']
    query_subcorpus = 'Results'
    analysis_client.cqp_query(corpus, query_subcorpus, query)
    """ Evaluate query results """
    match_corpus = '{}:{}'.format(corpus, query_subcorpus)
    match_num = min(int(message['hits_per_page']) - 1,
                    analysis_client.cqp_subcorpus_size(match_corpus))
    if match_num == 0:
        print('No matches found.')
        exit()
    if not analysis_client.cqp_subcorpus_has_field(match_corpus, CONST_FIELD_MATCH):
        print('Error.')
        exit()
    if not analysis_client.cqp_subcorpus_has_field(match_corpus, CONST_FIELD_MATCHEND):
        print('Error')
        exit()
    match_boundaries = zip(analysis_client.cqp_dump_subcorpus(match_corpus, CONST_FIELD_MATCH, 0, match_num - 1),
                           analysis_client.cqp_dump_subcorpus(match_corpus, CONST_FIELD_MATCHEND, 0, match_num - 1))
    matches = []
    for match_start, match_end in match_boundaries:
        matches.append({'cpos_list': list(range(match_start, match_end + 1))})
    cpos_list = []
    for match in matches:
        cpos_list = cpos_list + match['cpos_list']
    cpos_list = list(set(cpos_list))
    pos_list = analysis_client.cl_cpos2str('{}.pos'.format(corpus), cpos_list)
    word_list = analysis_client.cl_cpos2str('{}.word'.format(corpus), cpos_list)
    foo = {}
    for cpos, pos, word in zip(cpos_list, pos_list, word_list):
        foo[cpos] = {'pos': pos, 'word': word}
    for match in matches:
        match['pos_list'] = [foo[cpos]['pos'] for cpos in match['cpos_list']]
        match['word_list'] = [foo[cpos]['word'] for cpos in match['cpos_list']]
        match.pop('cpos_list', None)
    logger.warning(matches)
    socketio.emit('query', matches, room=request.sid)


def observe_corpus_analysis_connection(app, corpus_id, session_id):
    with app.app_context():
        corpus = Corpus.query.get(corpus_id)
        if corpus is None:
            ''' TODO: Handle this '''
            return
        while corpus.status != 'analysing':
            db.session.refresh(corpus)
            socketio.sleep(3)
        analysis_client = CQiClient(host='{}_analysis_container{}'.format(corpus.creator.username, corpus.id), port=4877)
        analysis_client.ctrl_connect('opaque', 'opaque')
        analysis_clients[session_id] = analysis_client
        socketio.emit('init_corpus_analysis', 'Ready', room=session_id)
        while session_id in connected_sessions:
            socketio.sleep(3)
        analysis_client.ctrl_bye()
        analysis_clients.pop(session_id, None)
        analysis_sessions[corpus_id].remove(session_id)
        if not analysis_sessions[corpus_id]:
            analysis_sessions.pop(corpus_id, None)
            corpus.status = 'stop analysis'
            db.session.commit()