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

import unittest
import tempfile
import os
from subprocess import check_call
from shutil import rmtree, copyfile
import urllib2
from urllib import quote
import json

# global settings for the test
# currently assumes a manually started server
solr = 'http://localhost:8080/solr/'
repo = tempfile.mkdtemp()
reponame = os.path.basename(repo)
repourl = 'file://' + repo
hook = repo + '/hooks/post-commit'
hooklog = hook + '.log'

def run(cmd):
  #not properly quoted:#print '# ' + ' '.join(cmd)
  return check_call(cmd)

def createRepository():
  ''' create repository with indexing hook enabled '''
  run(['svnadmin', 'create', repo])
  hook = repo + '/hooks/post-commit'
  f = open(hook, 'w')
  f.write('#!/bin/sh\n')
  f.write('export LC_ALL="en_US.UTF-8"\n')
  f.write('%s $1 $2' % os.path.abspath('../hook/svnhook.py'))
  f.write(' >> %s 2>&1\n' % hooklog)
  f.close()
  os.chmod(hook, 0777)

def createInitialStructure():
  '''
    Test data can be added from tests but we do a batch first for faster test execution.
    All tests may use this data.
  '''
  print '# ------- setup start --------'
  run(['svn', 'import', 'docs', repourl + '/docs', '-m', 'Mixed documents'])
  # properties must be added in a working copy, using a small subfolder for performance
  propwc = tempfile.mkdtemp()
  run(['svn', 'co', repourl + '/docs/svnprops/', propwc])
  txt =  propwc + '/textwithsvnprops.txt'
  run(['svn', 'propset', 'svn:keywords', 'Id LastChangedRevision HeadURL', txt])
  # Currently keywords properties are only treated as keywords if their namespace match the schema copyField rule
  run(['svn', 'propset', 'cms:keywords', 'ReposSearch repossearch keywordfromsvnprop', txt])
  run(['svn', 'propset', 'whatever:tags', 'metadataindexing tagging-in-svnprop', txt])
  run(['svn', 'propset', 'custom:someurl', 'Visit http://repossearch.com/', txt])
  run(['svn', 'propset', 'svn:mime-type', 'application/pdf', propwc + '/shortpdf.pdf'])
  run(['svn', 'propset', 'cms:keywords', 'keywordfromsvnprop', propwc + '/shortpdf.pdf'])
  # creating the invalid xml here so the error is not logged in rev 1 too
  xml = propwc + '/invalid, tag not closed.xml'
  f = open(xml, 'w')
  f.write('<?xml version="1.0"?>\n<document>\n<opentag>\n</document>\n')
  f.close()
  run(['svn', 'add', xml])
  run(['svn', 'propset', 'some:prop', 'OnInvalidXml', xml])
  run(['svn', 'commit', '-m', 'Properties added', propwc])
  print '# ---- common test data ----'
  run(['svnlook', 'tree', repo])
  print '# ------- hook log ---------'
  f = open(hooklog, 'r')
  print f.read()
  f.close()
  print '# ------- setup done --------'
  print ''

def curl(url):
  # transparently restrict hits to this instance
  if url.find('q=') > 0:
    url = url + '&fq=id_repo:' + reponame
  r = urllib2.urlopen(url)
  return r.read()

def search(query, queryType='standard'):
  '''
  queryType: standard/None, meta, content
  query: solr escaped but not urlencoded query
  '''
  # todo encode and stuff
  r = curl(solr + 'svnhead/select?qt=' + queryType + '&q=' + quote(query) + '&wt=json')
  #print '\n' + r
  return json.loads(r)

def searchMeta(query):
  return search(query, 'meta')

def searchContent(query):
  return search(query, 'content')

def s1(queryType, query):
  ''' search, throw exception if number if hits is not 1, return the document id for the hit '''
  r = search(query, queryType)
  n = r['response']['numFound']
  if n == 0:
    raise NameError('No results for query "%s" type %s' % (query, queryType))
  if n > 1:
    raise NameError('%d results for query "%s" type %s: \n%s' % (n, query, queryType, repr(r)))
  return r['response']['docs'][0]['id'].partition('^')[2]

class ReposSearchTest(unittest.TestCase):

  def testServerIsUp(self):
    r = curl(solr + '')
    self.assertTrue(r.find('svnhead') > -1, r)

  def testFilename(self):
    self.assertEqual(s1('meta', 'shouldBeUNIQUEfilename'),
                     '/docs/filenames/shouldBeUniqueFilename.txt')

  def testFilenameWithExtension(self):
    self.assertEqual(s1('meta', 'shouldbeuniquefilename.txt'),
                     '/docs/filenames/shouldBeUniqueFilename.txt')

  def testFilenameTokenizeDash(self):
    self.assertEqual(s1('standard', 'name:(subversion AND related AND document2)'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt',
                     'dash should be searchable as space')
    # How do we handle dashes when the adjacent chars are digits?
    self.assertEqual(s1('meta', 'subversion document 2010'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    self.assertEqual(s1('meta', '2010-02-10'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    # TODO Should dash-separated words be joined? Difference between digits and letters?
    #self.assertEqual(s1('meta', '20100210'),
    #                 '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    #self.assertEqual(s1('meta', 'subversionrelated 2010-02-10'),
    #                 '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    # TODO Should CamelCase be split into separate words? in that case we should have 2 hits above
    # Even if CamelCase is split it must still match as one word, as below
    self.assertEqual(s1('meta', 'subversionrelated ver'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt')

  def testFilenameTokenizerCodes(self):
    self.assertEqual(s1('meta', '11.2'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt',
                     'string like 11.2 (version number) must be searchable')
    self.assertEqual(s1('meta', 'DNR333'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt',
                     'letters+digits could be a product code and must be searchable')

  def testFilenameLatin1Accents(self):
    self.assertEqual(s1('meta', 'aeea'),
                     u'/docs/filenames/Latin1 accents áéèà.txt',
                     'accented chars should be searchable without accent')

  def testFilenameTokenizeWhitespaceComma(self):
    self.assertEqual(s1('standard', 'name:(short AND not AND long AND filename AND txt)'),
                     '/docs/filenames/short, not long, filename.txt',
                     'commas should be ignored, common words like "not" should be kept')

  def testFilenameUnicode(self):
    self.assertEqual(s1('meta', u'Swedish åäö'.encode('utf-8')),
                     u'/docs/filenames/In Swedish Åäö.txt')

  def testContentXslAndOds(self):
    r = searchContent('"cell B2"')
    self.assertEqual(r['response']['numFound'], 2)

  def testContentDocAndOdt(self):
    r = searchContent('"Repos Search is ok"')
    self.assertEqual(r['response']['numFound'], 2)

  def testPDFPrintedOnOsX(self):
    self.assertEqual(s1('content', 'veryshortpdfcontents'),
                     '/docs/svnprops/shortpdf.pdf')

  def testContentInvalidXml(self):
    r = searchMeta('invalid not closed')
    self.assertEqual(r['response']['numFound'], 1)
    error = r['response']['docs'][0]['text_error']
    self.assertTrue(error.find('xml') >= 0, 'Got: %s' % error)
    self.assertTrue(error.find('opentag') >= 0, 'Should specify the invalid tag. Got: %s' % error)
    self.assertEqual(r['response']['docs'][0]['svnprop_some_prop'], 'OnInvalidXml',
                     'svn properties should be indexed even if document can not be parsed')

  def testSvnPropsTextfile(self):
    r = searchMeta('textwithsvnprops.txt')
    self.assertEqual(r['response']['numFound'], 1)
    self.assertEqual(r['response']['docs'][0]['svnrevision'], 2)
    self.assertFalse(r['response']['docs'][0].has_key('svnprop_svn_keywords'),
                     'svn properties should normally not be indexed');
    self.assertEqual(r['response']['docs'][0]['svnprop_custom_someurl'], 'Visit http://repossearch.com/',
                     'by default custom properties should be indexed and stored as plain strings')

  def testMetaDontMatchSvnKeywords(self):
    r = searchMeta('LastChangedRevision')
    self.assertEqual(r['response']['numFound'], 0,
                     'svn keywords should not be indexed because they are not keywords')

  def testSearchOnCustomKeywords(self):
    r = searchMeta('keywordfromsvnprop')
    docs = r['response']['docs']
    self.assertEqual(len(docs), 2, 'should be two files with this keyword')
    self.assertEqual(r['response']['docs'][1]['id'], # could probably be index 0, ranking not tested
                     '%s^/docs/svnprops/textwithsvnprops.txt' % reponame,
                     'meta search should include keywords properties with likely namespaces except svn')

  def testSearchOnCustomTagsField(self):
    self.assertEqual(s1('meta', 'tagging-in-svnprop'),
                     '/docs/svnprops/textwithsvnprops.txt',
                     'meta search should include any *:kewords property value except svn:keywords')

  def testKeywordsFromPDFAndSvnProperty(self):
    r = search('allkeywords:(keywordinsaveaspdf AND keywordfromsvnprop)')
    self.assertEqual(r['response']['numFound'], 1, 'should match on both keyword from file and svn')
    doc = r['response']['docs'][0]
    self.assertEqual(doc['id'], '%s^/docs/svnprops/shortpdf.pdf' % reponame)
    self.assertEqual(doc['svnprop_svn_mime_type'], 'application/pdf', 'should index svn:mime-type')

  def testSearchEmbeddedPDFTitle(self):
    self.assertEqual(s1('meta', '"PDF Title"'),
                     '/docs/svnprops/shortpdf.pdf',
                     'should match on part of PDF embedded Title')

  def testSearchEmbeddedPDFSubject(self):
    # TODO Should subject be matched in meta query?
    self.assertEqual(s1('standard', 'subject:"PDF subject"'),
                     '/docs/svnprops/shortpdf.pdf',
                     'should match on part of PDF embedded Subject')

  def testContentUTF8Txt(self):
    self.assertEqual(s1('content', u'ÅÄÖåäö'.encode('utf-8')),
                     u'/docs/filenames/In Swedish Åäö.txt')

if __name__ == '__main__':
  createRepository()
  createInitialStructure()
  unittest.main()
  rmtree(repo)<commit_msg>Added test for filename starting with whitespace. I don't know any sensible use cases for such a name but anyway the hook should not crash.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import unittest
import tempfile
import os
from subprocess import check_call
from shutil import rmtree, copyfile
import urllib2
from urllib import quote
import json

# global settings for the test
# currently assumes a manually started server
solr = 'http://localhost:8080/solr/'
repo = tempfile.mkdtemp()
reponame = os.path.basename(repo)
repourl = 'file://' + repo
hook = repo + '/hooks/post-commit'
hooklog = hook + '.log'

def run(cmd):
  #not properly quoted:#print '# ' + ' '.join(cmd)
  return check_call(cmd)

def createRepository():
  ''' create repository with indexing hook enabled '''
  run(['svnadmin', 'create', repo])
  hook = repo + '/hooks/post-commit'
  f = open(hook, 'w')
  f.write('#!/bin/sh\n')
  f.write('export LC_ALL="en_US.UTF-8"\n')
  f.write('%s $1 $2' % os.path.abspath('../hook/svnhook.py'))
  f.write(' >> %s 2>&1\n' % hooklog)
  f.close()
  os.chmod(hook, 0777)

def createInitialStructure():
  '''
    Test data can be added from tests but we do a batch first for faster test execution.
    All tests may use this data.
  '''
  print '# ------- setup start --------'
  run(['svn', 'import', 'docs', repourl + '/docs', '-m', 'Mixed documents'])
  # properties must be added in a working copy, using a small subfolder for performance
  propwc = tempfile.mkdtemp()
  run(['svn', 'co', repourl + '/docs/svnprops/', propwc])
  txt =  propwc + '/textwithsvnprops.txt'
  run(['svn', 'propset', 'svn:keywords', 'Id LastChangedRevision HeadURL', txt])
  # Currently keywords properties are only treated as keywords if their namespace match the schema copyField rule
  run(['svn', 'propset', 'cms:keywords', 'ReposSearch repossearch keywordfromsvnprop', txt])
  run(['svn', 'propset', 'whatever:tags', 'metadataindexing tagging-in-svnprop', txt])
  run(['svn', 'propset', 'custom:someurl', 'Visit http://repossearch.com/', txt])
  run(['svn', 'propset', 'svn:mime-type', 'application/pdf', propwc + '/shortpdf.pdf'])
  run(['svn', 'propset', 'cms:keywords', 'keywordfromsvnprop', propwc + '/shortpdf.pdf'])
  # creating the invalid xml here so the error is not logged in rev 1 too
  xml = propwc + '/invalid, tag not closed.xml'
  f = open(xml, 'w')
  f.write('<?xml version="1.0"?>\n<document>\n<opentag>\n</document>\n')
  f.close()
  run(['svn', 'add', xml])
  run(['svn', 'propset', 'some:prop', 'OnInvalidXml', xml])
  run(['svn', 'commit', '-m', 'Properties added', propwc])
  print '# ---- common test data ----'
  run(['svnlook', 'tree', repo])
  print '# ------- hook log ---------'
  f = open(hooklog, 'r')
  print f.read()
  f.close()
  print '# ------- setup done --------'
  print ''

def curl(url):
  # transparently restrict hits to this instance
  if url.find('q=') > 0:
    url = url + '&fq=id_repo:' + reponame
  r = urllib2.urlopen(url)
  return r.read()

def search(query, queryType='standard'):
  '''
  queryType: standard/None, meta, content
  query: solr escaped but not urlencoded query
  '''
  # todo encode and stuff
  r = curl(solr + 'svnhead/select?qt=' + queryType + '&q=' + quote(query) + '&wt=json')
  #print '\n' + r
  return json.loads(r)

def searchMeta(query):
  return search(query, 'meta')

def searchContent(query):
  return search(query, 'content')

def s1(queryType, query):
  ''' search, throw exception if number if hits is not 1, return the document id for the hit '''
  r = search(query, queryType)
  n = r['response']['numFound']
  if n == 0:
    raise NameError('No results for query "%s" type %s' % (query, queryType))
  if n > 1:
    raise NameError('%d results for query "%s" type %s: \n%s' % (n, query, queryType, repr(r)))
  return r['response']['docs'][0]['id'].partition('^')[2]

class ReposSearchTest(unittest.TestCase):

  def testServerIsUp(self):
    r = curl(solr + '')
    self.assertTrue(r.find('svnhead') > -1, r)

  def testFilename(self):
    self.assertEqual(s1('meta', 'shouldBeUNIQUEfilename'),
                     '/docs/filenames/shouldBeUniqueFilename.txt')

  def testFilenameWithExtension(self):
    self.assertEqual(s1('meta', 'shouldbeuniquefilename.txt'),
                     '/docs/filenames/shouldBeUniqueFilename.txt')

  def testFilenameTokenizeDash(self):
    self.assertEqual(s1('standard', 'name:(subversion AND related AND document2)'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt',
                     'dash should be searchable as space')
    # How do we handle dashes when the adjacent chars are digits?
    self.assertEqual(s1('meta', 'subversion document 2010'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    self.assertEqual(s1('meta', '2010-02-10'),
                     '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    # TODO Should dash-separated words be joined? Difference between digits and letters?
    #self.assertEqual(s1('meta', '20100210'),
    #                 '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    #self.assertEqual(s1('meta', 'subversionrelated 2010-02-10'),
    #                 '/docs/filenames/Subversion-related Document2, 2010-02-10.txt')
    # TODO Should CamelCase be split into separate words? in that case we should have 2 hits above
    # Even if CamelCase is split it must still match as one word, as below
    self.assertEqual(s1('meta', 'subversionrelated ver'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt')

  def testFilenameTokenizerCodes(self):
    self.assertEqual(s1('meta', '11.2'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt',
                     'string like 11.2 (version number) must be searchable')
    self.assertEqual(s1('meta', 'DNR333'),
                     '/docs/filenames/SubversionRelated DNR333 ver 11.2.txt',
                     'letters+digits could be a product code and must be searchable')

  def testFilenameLatin1Accents(self):
    self.assertEqual(s1('meta', 'aeea'),
                     u'/docs/filenames/Latin1 accents áéèà.txt',
                     'accented chars should be searchable without accent')

  def testFilenameTokenizeWhitespaceComma(self):
    self.assertEqual(s1('standard', 'name:(short AND not AND long AND filename AND txt)'),
                     '/docs/filenames/short, not long, filename.txt',
                     'commas should be ignored, common words like "not" should be kept')

  def testFilenameUnicode(self):
    self.assertEqual(s1('meta', u'Swedish åäö'.encode('utf-8')),
                     u'/docs/filenames/In Swedish Åäö.txt')

  def testContentXslAndOds(self):
    r = searchContent('"cell B2"')
    self.assertEqual(r['response']['numFound'], 2)

  def testContentDocAndOdt(self):
    r = searchContent('"Repos Search is ok"')
    self.assertEqual(r['response']['numFound'], 2)

  def testPDFPrintedOnOsX(self):
    self.assertEqual(s1('content', 'veryshortpdfcontents'),
                     '/docs/svnprops/shortpdf.pdf')

  def testContentInvalidXml(self):
    r = searchMeta('invalid not closed')
    self.assertEqual(r['response']['numFound'], 1)
    error = r['response']['docs'][0]['text_error']
    self.assertTrue(error.find('xml') >= 0, 'Got: %s' % error)
    self.assertTrue(error.find('opentag') >= 0, 'Should specify the invalid tag. Got: %s' % error)
    self.assertEqual(r['response']['docs'][0]['svnprop_some_prop'], 'OnInvalidXml',
                     'svn properties should be indexed even if document can not be parsed')

  def testSvnPropsTextfile(self):
    r = searchMeta('textwithsvnprops.txt')
    self.assertEqual(r['response']['numFound'], 1)
    self.assertEqual(r['response']['docs'][0]['svnrevision'], 2)
    self.assertFalse(r['response']['docs'][0].has_key('svnprop_svn_keywords'),
                     'svn properties should normally not be indexed');
    self.assertEqual(r['response']['docs'][0]['svnprop_custom_someurl'], 'Visit http://repossearch.com/',
                     'by default custom properties should be indexed and stored as plain strings')

  def testMetaDontMatchSvnKeywords(self):
    r = searchMeta('LastChangedRevision')
    self.assertEqual(r['response']['numFound'], 0,
                     'svn keywords should not be indexed because they are not keywords')

  def testSearchOnCustomKeywords(self):
    r = searchMeta('keywordfromsvnprop')
    docs = r['response']['docs']
    self.assertEqual(len(docs), 2, 'should be two files with this keyword')
    self.assertEqual(r['response']['docs'][1]['id'], # could probably be index 0, ranking not tested
                     '%s^/docs/svnprops/textwithsvnprops.txt' % reponame,
                     'meta search should include keywords properties with likely namespaces except svn')

  def testSearchOnCustomTagsField(self):
    self.assertEqual(s1('meta', 'tagging-in-svnprop'),
                     '/docs/svnprops/textwithsvnprops.txt',
                     'meta search should include any *:kewords property value except svn:keywords')

  def testKeywordsFromPDFAndSvnProperty(self):
    r = search('allkeywords:(keywordinsaveaspdf AND keywordfromsvnprop)')
    self.assertEqual(r['response']['numFound'], 1, 'should match on both keyword from file and svn')
    doc = r['response']['docs'][0]
    self.assertEqual(doc['id'], '%s^/docs/svnprops/shortpdf.pdf' % reponame)
    self.assertEqual(doc['svnprop_svn_mime_type'], 'application/pdf', 'should index svn:mime-type')

  def testSearchEmbeddedPDFTitle(self):
    self.assertEqual(s1('meta', '"PDF Title"'),
                     '/docs/svnprops/shortpdf.pdf',
                     'should match on part of PDF embedded Title')

  def testSearchEmbeddedPDFSubject(self):
    # TODO Should subject be matched in meta query?
    self.assertEqual(s1('standard', 'subject:"PDF subject"'),
                     '/docs/svnprops/shortpdf.pdf',
                     'should match on part of PDF embedded Subject')

  def testContentUTF8Txt(self):
    self.assertEqual(s1('content', u'ÅÄÖåäö'.encode('utf-8')),
                     u'/docs/filenames/In Swedish Åäö.txt')

  def testFilenameWithLeadingWhitespace(self):
    '''Subversion and some OSes support filenames that start with whitespace'''
    (h, f) = tempfile.mkstemp()
    run(['svn', 'import', "%s" % f, repourl + '/ leadingwhitespace.txt', '-m', 'OK name?'])
    os.remove(f)
    self.assertEqual(s1('meta', 'leadingwhitespace'),
                     '/ leadingwhitespace.txt',
                     'Subversion allowes names taht start with whitespace so we should too')

if __name__ == '__main__':
  createRepository()
  createInitialStructure()
  unittest.main()
  rmtree(repo)<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Converts sentence.doc_char_begin, sentence.doc_char_end to sentence.tokens.
"""
import sys
from tqdm import tqdm
from collections import namedtuple
from kbpo import db

def do_command(_):
    with db.CONN:
        with db.CONN.cursor() as cur:
            # Create temporary table to replace sentence.
            cur.execute("""DROP TABLE IF EXISTS _sentence;""")
            cur.execute("""
                CREATE TABLE _sentence (
                  id INTEGER NOT NULL DEFAULT nextval('sentence_id_seq'),
                  updated TIMESTAMP NOT NULL DEFAULT (now() at time zone 'utc'),

                  doc_id TEXT NOT NULL,
                  span INT4RANGE NOT NULL,
                  sentence_index SMALLINT NOT NULL,

                  gloss TEXT NOT NULL,

                  token_spans INT4RANGE[] NOT NULL,
                  words TEXT[] NOT NULL,
                  lemmas TEXT[] NOT NULL,
                  pos_tags TEXT[] NOT NULL,
                  ner_tags TEXT[] NOT NULL,
                  dependencies TEXT NOT NULL
                );""")
            NewRow = namedtuple("NewRow", "id updated doc_id span sentence_index gloss tokens words lemmas pos_tags ner_tags dependencies".split())

            values = []
            cur.execute("""SELECT * FROM sentence""")
            for row in tqdm(cur, total=cur.rowcount):
                begin, end = row.doc_char_begin, row.doc_char_end
                tokens = [db.Int4NumericRange(int(b), int(e)) for b, e in zip(begin, end)]
                row_ = NewRow(row.id, row.updated, row.doc_id, row.span, row.sentence_index, row.gloss, tokens, row.words, row.lemmas, row.pos_tags, row.ner_tags, row.dependencies)
                values.append(row_)
            db.execute_values(cur, """INSERT INTO _sentence(id, updated, doc_id, span, sentence_index, gloss, token_span, words, lemmas, pos_tags, ner_tags, dependencies) VALUES %s""", values)

            # ADD constraints
            cur.execute("""
                ALTER TABLE _sentence ADD CONSTRAINT _sentence_pkey PRIMARY KEY (doc_id, span);
                ALTER TABLE _sentence ADD CONSTRAINT _sentence_doc_id_fkey FOREIGN KEY (doc_id) REFERENCES document(id);
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_tokens_length CHECK ((array_lower(words, 1) = array_lower(token_spans, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_lemma_length CHECK ((array_lower(words, 1) = array_lower(lemmas, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_ner_length CHECK ((array_lower(words, 1) = array_lower(ner_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_pos_length CHECK ((array_lower(words, 1) = array_lower(pos_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_tokens_length CHECK ((array_upper(words, 1) = array_upper(token_spans, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_lemma_length CHECK ((array_upper(words, 1) = array_upper(lemmas, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_ner_length CHECK ((array_upper(words, 1) = array_upper(ner_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_pos_length CHECK ((array_upper(words, 1) = array_upper(pos_tags, 1)));
                """)
            cur.execute("""DROP TABLE sentence;""")
            cur.execute("""ALTER TABLE _sentence RENAME TO sentence;""")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Converts sentence.doc_char_begin, sentence.doc_char_end to sentence.tokens.')
    parser.set_defaults(func=do_command)

    ARGS = parser.parse_args()
    if ARGS.func is None:
        parser.print_help()
        sys.exit(1)
    else:
        ARGS.func(ARGS)
<commit_msg>Fixed update spans script<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Converts sentence.doc_char_begin, sentence.doc_char_end to sentence.tokens.
"""
import sys
from tqdm import tqdm
from collections import namedtuple
from kbpo import db

def do_command(_):
    with db.CONN:
        with db.CONN.cursor() as cur:
            # Create temporary table to replace sentence.
            cur.execute("""DROP TABLE IF EXISTS _sentence;""")
            cur.execute("""
                CREATE TABLE _sentence (
                  id INTEGER NOT NULL DEFAULT nextval('sentence_id_seq'),
                  updated TIMESTAMP NOT NULL DEFAULT (now() at time zone 'utc'),

                  doc_id TEXT NOT NULL,
                  span INT4RANGE NOT NULL,
                  sentence_index SMALLINT NOT NULL,

                  gloss TEXT NOT NULL,

                  token_spans INT4RANGE[] NOT NULL,
                  words TEXT[] NOT NULL,
                  lemmas TEXT[] NOT NULL,
                  pos_tags TEXT[] NOT NULL,
                  ner_tags TEXT[] NOT NULL,
                  dependencies TEXT NOT NULL
                );""")
            NewRow = namedtuple("NewRow", "id updated doc_id span sentence_index gloss tokens words lemmas pos_tags ner_tags dependencies".split())

            values = []
            cur.execute("""SELECT * FROM sentence""")
            for row in tqdm(cur, total=cur.rowcount):
                begin, end = row.doc_char_begin, row.doc_char_end
                tokens = [db.Int4NumericRange(int(b), int(e)) for b, e in zip(begin, end)]
                row_ = NewRow(row.id, row.updated, row.doc_id, row.span, row.sentence_index, row.gloss, tokens, row.words, row.lemmas, row.pos_tags, row.ner_tags, row.dependencies)
                values.append(row_)
            db.execute_values(cur, """INSERT INTO _sentence(id, updated, doc_id, span, sentence_index, gloss, token_spans, words, lemmas, pos_tags, ner_tags, dependencies) VALUES %s""", values)

            # ADD constraints
            cur.execute("""
                ALTER TABLE _sentence ADD CONSTRAINT _sentence_pkey PRIMARY KEY (doc_id, span);
                ALTER TABLE _sentence ADD CONSTRAINT _sentence_doc_id_fkey FOREIGN KEY (doc_id) REFERENCES document(id);
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_tokens_length CHECK ((array_lower(words, 1) = array_lower(token_spans, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_lemma_length CHECK ((array_lower(words, 1) = array_lower(lemmas, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_ner_length CHECK ((array_lower(words, 1) = array_lower(ner_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT __word_length_same_as_pos_length CHECK ((array_lower(words, 1) = array_lower(pos_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_tokens_length CHECK ((array_upper(words, 1) = array_upper(token_spans, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_lemma_length CHECK ((array_upper(words, 1) = array_upper(lemmas, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_ner_length CHECK ((array_upper(words, 1) = array_upper(ner_tags, 1)));
                ALTER TABLE _sentence ADD CONSTRAINT _word_length_same_as_pos_length CHECK ((array_upper(words, 1) = array_upper(pos_tags, 1)));
                """)
            cur.execute("""DROP TABLE sentence;""")
            cur.execute("""ALTER TABLE _sentence RENAME TO sentence;""")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Converts sentence.doc_char_begin, sentence.doc_char_end to sentence.tokens.')
    parser.set_defaults(func=do_command)

    ARGS = parser.parse_args()
    if ARGS.func is None:
        parser.print_help()
        sys.exit(1)
    else:
        ARGS.func(ARGS)
<|endoftext|>


<commit_before>'''
Created on Aug 6, 2015

@author: Connor
'''
from os import path

from sphinxbase.sphinxbase import Config, Config_swigregister #@UnusedImport
from pocketsphinx.pocketsphinx import Decoder
import pyaudio, speech_recognition
import client.tts as tts

MODEL_DIR = "..\models"

# Must be in the sphinx dict file
WAKE_UP_WORD = "athena"
ERROR_MESSAGE = "Sorry, I could not understand that."

def init():
    # Create a decoder with certain model
    config = Decoder.default_config()
    config.set_string('-logfn', '..\logs\passive-listen.log')
    config.set_string('-hmm', path.join(MODEL_DIR, 'en-us\en-us'))
    config.set_string('-lm', path.join(MODEL_DIR, 'en-us\en-us.lm.dmp'))
    config.set_string('-dict', path.join(MODEL_DIR, 'en-us\cmudict-en-us.dict'))

    # Decode streaming data
    global decoder, p
    decoder = Decoder(config)
    decoder.set_keyphrase("wakeup", WAKE_UP_WORD)
    decoder.set_search("wakeup")

    p = pyaudio.PyAudio()

def listen_keyword():
    global decoder, p
    stream = p.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True, frames_per_buffer=1024)
    stream.start_stream()
    p.get_default_input_device_info()

    print("~ Passive listening... ")
    decoder.start_utt()
    while True:
        buf = stream.read(1024)
        decoder.process_raw(buf, False, False)
        if decoder.hyp() != None and decoder.hyp().hypstr == WAKE_UP_WORD:
            decoder.end_utt()
            return
        else:
            continue
    decoder.end_utt()

def active_listen():
    r = speech_recognition.Recognizer()

    with speech_recognition.Microphone() as src:    # use the default microphone as the audio source
        r.adjust_for_ambient_noise(src)             # listen for 1 second to calibrate the energy threshold for ambient noise levels
        print("\n~ Active listening... ")
        tts.play_mp3("double-beep.mp3")
        audio = r.listen(src)                       # listen for the first phrase and extract it into audio data

    try:
        msg = r.recognize(audio)                    # recognize speech using Google Speech Recognition
        print("\n~ \""+msg+"\"")
    except LookupError:                             # speech is unintelligible
        msg = ""
        print("\n~ "+ERROR_MESSAGE+"\n")
        tts.speak(ERROR_MESSAGE)
    finally:
        return msg
<commit_msg>Minor STT update<commit_after>'''
Created on Aug 6, 2015

@author: Connor
'''
from os import path

from sphinxbase.sphinxbase import Config, Config_swigregister #@UnusedImport
from pocketsphinx.pocketsphinx import Decoder
import pyaudio, speech_recognition
import client.tts as tts
import os

MODEL_DIR = "..\models"
LOGS_DIR = "..\logs"

# Must be in the sphinx dict file
WAKE_UP_WORD = "athena"
ERROR_MESSAGE = "Sorry, I could not understand that."

def init():
    # Be wary of an OSError due to a race condition
    if not os.path.exists(LOGS_DIR):
        os.makedirs(LOGS_DIR)

    # Create a decoder with certain model
    config = Decoder.default_config()
    config.set_string('-logfn', path.join(LOGS_DIR, 'passive-listen.log'))
    config.set_string('-hmm', path.join(MODEL_DIR, 'en-us\en-us'))
    config.set_string('-lm', path.join(MODEL_DIR, 'en-us\en-us.lm.dmp'))
    config.set_string('-dict', path.join(MODEL_DIR, 'en-us\cmudict-en-us.dict'))

    # Decode streaming data
    global decoder, p
    decoder = Decoder(config)
    decoder.set_keyphrase("wakeup", WAKE_UP_WORD)
    decoder.set_search("wakeup")

    p = pyaudio.PyAudio()

def listen_keyword():
    global decoder, p
    stream = p.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True, frames_per_buffer=1024)
    stream.start_stream()
    p.get_default_input_device_info()

    print("~ Passive listening... ")
    decoder.start_utt()
    while True:
        buf = stream.read(1024)
        decoder.process_raw(buf, False, False)
        if decoder.hyp() != None and decoder.hyp().hypstr == WAKE_UP_WORD:
            decoder.end_utt()
            return
        else:
            continue
    decoder.end_utt()

def active_listen():
    r = speech_recognition.Recognizer()

    with speech_recognition.Microphone() as src:    # use the default microphone as the audio source
        r.adjust_for_ambient_noise(src)             # listen for 1 second to calibrate the energy threshold for ambient noise levels
        print("\n~ Active listening... ")
        tts.play_mp3("double-beep.mp3")
        audio = r.listen(src)                       # listen for the first phrase and extract it into audio data

    try:
        msg = r.recognize(audio)                    # recognize speech using Google Speech Recognition
        print("\n~ \""+msg+"\"")
    except LookupError:                             # speech is unintelligible
        msg = ""
        print("\n~ "+ERROR_MESSAGE+"\n")
        tts.speak(ERROR_MESSAGE)
    finally:
        return msg
<|endoftext|>


<commit_before>import functools
import itertools
from collections import defaultdict
import time
import logging as log
from Queue import Queue, Empty
import weakref
import threading

from pykafka import base
from pykafka.common import OffsetType
from pykafka.exceptions import (OffsetOutOfRangeError, UnknownTopicOrPartition,
                                InvalidMessageError)

from .utils.error_handlers import handle_partition_responses, raise_error
from .protocol import (PartitionFetchRequest, PartitionOffsetCommitRequest,
                       PartitionOffsetFetchRequest, PartitionOffsetRequest)

class SimpleConsumer(base.BaseSimpleConsumer):
    def __init__(self,
                 topic,
                 cluster,
                 consumer_group=None,
                 partitions=None,
                 fetch_message_max_bytes=1024 * 1024,
                 num_consumer_fetchers=1,
                 auto_commit_enable=False,
                 auto_commit_interval_ms=60 * 1000,
                 queued_max_messages=2000,
                 fetch_min_bytes=1,
                 fetch_wait_max_ms=100,
                 refresh_leader_backoff_ms=200,
                 offsets_channel_backoff_ms=1000,
                 offsets_commit_max_retries=5,
                 auto_offset_reset=OffsetType.LATEST,
                 consumer_timeout_ms=-1):
        """Create a SimpleConsumer.

        Settings and default values are taken from the Scala
        consumer implementation.  Consumer group is included
        because it's necessary for offset management, but doesn't imply
        that this is a balancing consumer. Use a BalancedConsumer for
        that.

        :param topic: the topic this consumer should consume
        :type topic: pykafka.topic.Topic
        :param cluster: the cluster this consumer should connect to
        :type cluster: pykafka.cluster.Cluster
        :param consumer_group: the name of the consumer group to join
        :type consumer_group: str
        :param partitions: existing partitions to which to connect
        :type partitions: list of pykafka.partition.Partition
        :param fetch_message_max_bytes: the number of bytes of messages to
            attempt to fetch
        :type fetch_message_max_bytes: int
        :param num_consumer_fetchers: the number of threads used to fetch data
        :type num_consumer_fetchers: int
        :param auto_commit_enable: if true, periodically commit to kafka the
            offset of messages already fetched by this consumer
        :type auto_commit_enable: bool
        :param auto_commit_interval_ms: the frequency in ms that the consumer
            offsets are committed to kafka
        :type auto_commit_interval_ms: int
        :param queued_max_messages: max number of messages buffered for
            consumption
        :type queued_max_messages: int
        :param fetch_min_bytes: the minimum amount of data the server should
            return for a fetch request. If insufficient data is available the
            request will block
        :type fetch_min_bytes: int
        :param fetch_wait_max_ms: the maximum amount of time the server will
            block before answering the fetch request if there isn't sufficient
            data to immediately satisfy fetch_min_bytes
        :type fetch_wait_max_ms: int
        :param refresh_leader_backoff_ms: backoff time to refresh the leader of
            a partition after it loses the current leader
        :type refresh_leader_backoff_ms: int
        :param offsets_channel_backoff_ms: backoff time to retry offset
            commits/fetches
        :type offsets_channel_backoff_ms: int
        :param offsets_commit_max_retries: Retry the offset commit up to this
            many times on failure.
        :type offsets_commit_max_retries: int
        :param auto_offset_reset: what to do if an offset is out of range
        :type auto_offset_reset: int
        :param consumer_timeout_ms: throw a timeout exception to the consumer
            if no message is available for consumption after the specified interval
        :type consumer_timeout_ms: int
        """
        if not isinstance(cluster, weakref.ProxyType) and \
                not isinstance(cluster, weakref.CallableProxyType):
            self._cluster = weakref.proxy(cluster)
        else:
            self._cluster = cluster
        self._consumer_group = consumer_group
        self._topic = topic
        self._fetch_message_max_bytes = fetch_message_max_bytes
        self._fetch_min_bytes = fetch_min_bytes
        self._queued_max_messages = queued_max_messages
        self._num_consumer_fetchers = num_consumer_fetchers
        self._fetch_wait_max_ms = fetch_wait_max_ms
        self._consumer_timeout_ms = consumer_timeout_ms
        self._offsets_channel_backoff_ms = offsets_channel_backoff_ms
        self._offsets_commit_max_retries = offsets_commit_max_retries
        self._auto_offset_reset = auto_offset_reset

        self._last_message_time = time.time()

        self._auto_commit_enable = auto_commit_enable
        self._auto_commit_interval_ms = auto_commit_interval_ms
        self._last_auto_commit = time.time()

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        owned_partition_partial = functools.partial(
            OwnedPartition, consumer_group=self._consumer_group)
        if partitions:
            self._partitions = {owned_partition_partial(p): p
                                for p in partitions}
        else:
            self._partitions = {owned_partition_partial(p): topic.partitions[k]
                                for k, p in topic.partitions.iteritems()}
        self._partitions_by_id = {p.partition.id: p
                                  for p in self._partitions.iterkeys()}
        # Organize partitions by leader for efficient queries
        self._partitions_by_leader = defaultdict(list)
        for p in self._partitions.iterkeys():
            self._partitions_by_leader[p.partition.leader].append(p)
        self.partition_cycle = itertools.cycle(self._partitions.keys())

        self._default_error_handlers = self._build_default_error_handlers()

        self._running = True

        if self._auto_commit_enable:
            self._autocommit_worker_thread = self._setup_autocommit_worker()
        # we need to get the most up-to-date offsets before starting consumption
        self.fetch_offsets()
        self._fetch_workers = self._setup_fetch_workers()

    def _build_default_error_handlers(self):
        def _handle_UnknownTopicOrPartition(parts):
            raise_error(UnknownTopicOrPartition)

        def _handle_OffsetOutOfRangeError(parts):
            self._reset_offsets((owned_partition
                                 for owned_partition, pres in parts))

        def _handle_InvalidMessageError(parts):
            raise_error(InvalidMessageError)

        return {
            UnknownTopicOrPartition.ERROR_CODE: _handle_UnknownTopicOrPartition,
            OffsetOutOfRangeError.ERROR_CODE: _handle_OffsetOutOfRangeError,
            InvalidMessageError.ERROR_CODE: _handle_InvalidMessageError
        }

    @property
    def topic(self):
        return self._topic

    @property
    def partitions(self):
        return self._partitions

    def __del__(self):
        self.stop()

    def stop(self):
        self._running = False

    def _setup_autocommit_worker(self):
        def autocommitter():
            while True:
                if not self._running:
                    break
                if self._auto_commit_enable:
                    self._auto_commit()
        log.debug("Starting autocommitter thread")
        return self._cluster.handler.spawn(autocommitter)

    def _setup_fetch_workers(self):
        def fetcher():
            while True:
                if not self._running:
                    break
                self.fetch()
        return [self._cluster.handler.spawn(fetcher)
                for i in xrange(self._num_consumer_fetchers)]

    def __iter__(self):
        while True:
            message = self.consume()
            if not message and self._consumer_timed_out():
                raise StopIteration
            yield message

    def _consumer_timed_out(self):
        if self._consumer_timeout_ms == -1:
            return False
        disp = (time.time() - self._last_message_time) * 1000.0
        return disp > self._consumer_timeout_ms

    def consume(self):
        """Get one message from the consumer.

        :param timeout: Seconds to wait before returning None
        """
        owned_partition = self.partition_cycle.next()
        message = owned_partition.consume()

        if message:
            self._last_message_time = time.time()
            return message

    def _auto_commit(self):
        if not self._auto_commit_enable or self._auto_commit_interval_ms == 0:
            return

        if (time.time() - self._last_auto_commit) * 1000.0 >= self._auto_commit_interval_ms:
            log.info("Autocommitting consumer offset for consumer group %s and topic %s",
                     self._consumer_group, self._topic.name)
            self.commit_offsets()
            self._last_auto_commit = time.time()

    def commit_offsets(self):
        """Commit offsets for this consumer's topic

        Uses the offset commit/fetch API
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to commit offsets")

        reqs = [p.build_offset_commit_request() for p in self._partitions.keys()]
        log.info("Committing offsets for %d partitions", len(reqs))
        for i in xrange(self._offsets_commit_max_retries):
            if i > 0:
                log.debug("Retrying")
            time.sleep(i * (self._offsets_channel_backoff_ms / 1000))

            response = self._offset_manager.commit_consumer_group_offsets(
                self._consumer_group, 1, 'pykafka', reqs)
            parts_by_error = handle_partition_responses(
                response,
                self._default_error_handlers,
                partitions_by_id=self._partitions_by_id)
            if len(parts_by_error) == 1 and 0 in parts_by_error:
                break
            log.error("Error committing offsets for topic %s", self._topic.name)

            # retry only the partitions that errored
            parts_by_error.pop(0)
            errored_partitions = [op for err_group in parts_by_error.iteritems() for op in err_group]
            reqs = [p.build_offset_commit_request() for p in errored_partitions]

    def fetch_offsets(self):
        """Fetch offsets for this consumer's topic

        Uses the offset commit/fetch API
        Should be called when consumer starts and after any errors
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to fetch offsets")

        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.set_offset(pres.offset)

        log.info("Fetching offsets")

        reqs = [p.build_offset_fetch_request() for p in self._partitions.keys()]
        res = self._offset_manager.fetch_consumer_group_offsets(self._consumer_group, reqs)
        handle_partition_responses(
            res,
            self._default_error_handlers,
            success_handler=_handle_success,
            partitions_by_id=self._partitions_by_id)

    def _reset_offsets(self, errored_partitions):
        """Reset offsets after an OffsetOutOfRangeError

        Issue an OffsetRequest for each partition and set the appropriate
        returned offset in the OwnedPartition per self._auto_offset_reset

        :param errored_partitions: the partitions with out-of-range offsets
        :type errored_partitions: Iterable of OwnedPartition
        """
        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.set_offset(pres.offset[0])

        # group out-of-range partitions by leader
        by_leader = defaultdict(list)
        for p in errored_partitions:
            by_leader[p.partition.leader].append(p)

        # get valid offset ranges for each partition
        for broker, owned_partitions in by_leader.iteritems():
            reqs = [owned_partition.build_offset_request(self._auto_offset_reset)
                    for owned_partition in owned_partitions]
            response = broker.request_offset_limits(reqs)
            handle_partition_responses(
                response,
                self._default_error_handlers,
                success_handler=_handle_success,
                partitions_by_id=self._partitions_by_id)

    def fetch(self):
        """Fetch new messages for all partitions

        Create a FetchRequest for each broker and send it. Enqueue each of the
        returned messages in the approprate OwnedPartition.
        """
        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.enqueue_messages(pres.messages)

        for broker, owned_partitions in self._partitions_by_leader.iteritems():
            partition_reqs = []
            for owned_partition in owned_partitions:
                # attempt to acquire lock, just pass if we can't
                if owned_partition.lock.acquire(False) and \
                        owned_partition.message_count < self._queued_max_messages:
                    fetch_req = owned_partition.build_fetch_request(
                        self._fetch_message_max_bytes)
                    partition_reqs.append((owned_partition, fetch_req))
            if partition_reqs:
                response = broker.fetch_messages(
                    [a[1] for a in partition_reqs],
                    timeout=self._fetch_wait_max_ms,
                    min_bytes=self._fetch_min_bytes
                )
                handle_partition_responses(
                    response,
                    self._default_error_handlers,
                    success_handler=_handle_success,
                    partitions_by_id=self._partitions_by_id)
            for owned_partition, _ in partition_reqs:
                owned_partition.lock.release()

class OwnedPartition(object):
    """A partition that is owned by a SimpleConsumer.

    Used to keep track of offsets and the internal message queue.
    """

    def __init__(self,
                 partition,
                 consumer_group=None):
        self.partition = partition
        self._messages = Queue()
        self.last_offset_consumed = 0
        self.next_offset = 0
        self.lock = threading.Lock()

    @property
    def message_count(self):
        return self._messages.qsize()

    def set_offset(self, last_offset_consumed):
        """Set the internal offset counters from an OffsetFetchResponse

        :param last_offset_consumed: the last committed offset for this
            partition
        :type last_offset_consumed: int
        """
        self.last_offset_consumed = last_offset_consumed
        self.next_offset = last_offset_consumed + 1

    def build_offset_request(self, auto_offset_reset):
        """Create a PartitionOffsetRequest for this partition
        """
        return PartitionOffsetRequest(
            self.partition.topic.name, self.partition.id,
            auto_offset_reset, 1)

    def build_fetch_request(self, max_bytes):
        """Create a FetchPartitionRequest for this partition
        """
        return PartitionFetchRequest(
            self.partition.topic.name, self.partition.id,
            self.next_offset, max_bytes)

    def build_offset_commit_request(self):
        """Create a PartitionOffsetCommitRequest for this partition
        """
        return PartitionOffsetCommitRequest(
            self.partition.topic.name,
            self.partition.id,
            self.last_offset_consumed,
            int(time.time()),
            'pykafka'
        )

    def build_offset_fetch_request(self):
        """Create a PartitionOffsetFetchRequest for this partition
        """
        return PartitionOffsetFetchRequest(
            self.partition.topic.name,
            self.partition.id
        )

    def consume(self):
        """Get a single message from this partition
        """
        try:
            message = self._messages.get_nowait()
            self.last_offset_consumed = message.offset
            return message
        except Empty:
            return None

    def enqueue_messages(self, messages):
        for message in messages:
            if message.offset < self.last_offset_consumed:
                continue
            self._messages.put(message)
            self.next_offset = message.offset + 1
<commit_msg>simplify redundant handler definitions<commit_after>import functools
import itertools
from collections import defaultdict
import time
import logging as log
from Queue import Queue, Empty
import weakref
import threading

from pykafka import base
from pykafka.common import OffsetType
from pykafka.exceptions import (OffsetOutOfRangeError, UnknownTopicOrPartition,
                                InvalidMessageError)

from .utils.error_handlers import handle_partition_responses, raise_error
from .protocol import (PartitionFetchRequest, PartitionOffsetCommitRequest,
                       PartitionOffsetFetchRequest, PartitionOffsetRequest)

class SimpleConsumer(base.BaseSimpleConsumer):
    def __init__(self,
                 topic,
                 cluster,
                 consumer_group=None,
                 partitions=None,
                 fetch_message_max_bytes=1024 * 1024,
                 num_consumer_fetchers=1,
                 auto_commit_enable=False,
                 auto_commit_interval_ms=60 * 1000,
                 queued_max_messages=2000,
                 fetch_min_bytes=1,
                 fetch_wait_max_ms=100,
                 refresh_leader_backoff_ms=200,
                 offsets_channel_backoff_ms=1000,
                 offsets_commit_max_retries=5,
                 auto_offset_reset=OffsetType.LATEST,
                 consumer_timeout_ms=-1):
        """Create a SimpleConsumer.

        Settings and default values are taken from the Scala
        consumer implementation.  Consumer group is included
        because it's necessary for offset management, but doesn't imply
        that this is a balancing consumer. Use a BalancedConsumer for
        that.

        :param topic: the topic this consumer should consume
        :type topic: pykafka.topic.Topic
        :param cluster: the cluster this consumer should connect to
        :type cluster: pykafka.cluster.Cluster
        :param consumer_group: the name of the consumer group to join
        :type consumer_group: str
        :param partitions: existing partitions to which to connect
        :type partitions: list of pykafka.partition.Partition
        :param fetch_message_max_bytes: the number of bytes of messages to
            attempt to fetch
        :type fetch_message_max_bytes: int
        :param num_consumer_fetchers: the number of threads used to fetch data
        :type num_consumer_fetchers: int
        :param auto_commit_enable: if true, periodically commit to kafka the
            offset of messages already fetched by this consumer
        :type auto_commit_enable: bool
        :param auto_commit_interval_ms: the frequency in ms that the consumer
            offsets are committed to kafka
        :type auto_commit_interval_ms: int
        :param queued_max_messages: max number of messages buffered for
            consumption
        :type queued_max_messages: int
        :param fetch_min_bytes: the minimum amount of data the server should
            return for a fetch request. If insufficient data is available the
            request will block
        :type fetch_min_bytes: int
        :param fetch_wait_max_ms: the maximum amount of time the server will
            block before answering the fetch request if there isn't sufficient
            data to immediately satisfy fetch_min_bytes
        :type fetch_wait_max_ms: int
        :param refresh_leader_backoff_ms: backoff time to refresh the leader of
            a partition after it loses the current leader
        :type refresh_leader_backoff_ms: int
        :param offsets_channel_backoff_ms: backoff time to retry offset
            commits/fetches
        :type offsets_channel_backoff_ms: int
        :param offsets_commit_max_retries: Retry the offset commit up to this
            many times on failure.
        :type offsets_commit_max_retries: int
        :param auto_offset_reset: what to do if an offset is out of range
        :type auto_offset_reset: int
        :param consumer_timeout_ms: throw a timeout exception to the consumer
            if no message is available for consumption after the specified interval
        :type consumer_timeout_ms: int
        """
        if not isinstance(cluster, weakref.ProxyType) and \
                not isinstance(cluster, weakref.CallableProxyType):
            self._cluster = weakref.proxy(cluster)
        else:
            self._cluster = cluster
        self._consumer_group = consumer_group
        self._topic = topic
        self._fetch_message_max_bytes = fetch_message_max_bytes
        self._fetch_min_bytes = fetch_min_bytes
        self._queued_max_messages = queued_max_messages
        self._num_consumer_fetchers = num_consumer_fetchers
        self._fetch_wait_max_ms = fetch_wait_max_ms
        self._consumer_timeout_ms = consumer_timeout_ms
        self._offsets_channel_backoff_ms = offsets_channel_backoff_ms
        self._offsets_commit_max_retries = offsets_commit_max_retries
        self._auto_offset_reset = auto_offset_reset

        self._last_message_time = time.time()

        self._auto_commit_enable = auto_commit_enable
        self._auto_commit_interval_ms = auto_commit_interval_ms
        self._last_auto_commit = time.time()

        self._offset_manager = self._cluster.get_offset_manager(self._consumer_group)

        owned_partition_partial = functools.partial(
            OwnedPartition, consumer_group=self._consumer_group)
        if partitions:
            self._partitions = {owned_partition_partial(p): p
                                for p in partitions}
        else:
            self._partitions = {owned_partition_partial(p): topic.partitions[k]
                                for k, p in topic.partitions.iteritems()}
        self._partitions_by_id = {p.partition.id: p
                                  for p in self._partitions.iterkeys()}
        # Organize partitions by leader for efficient queries
        self._partitions_by_leader = defaultdict(list)
        for p in self._partitions.iterkeys():
            self._partitions_by_leader[p.partition.leader].append(p)
        self.partition_cycle = itertools.cycle(self._partitions.keys())

        self._default_error_handlers = self._build_default_error_handlers()

        self._running = True

        if self._auto_commit_enable:
            self._autocommit_worker_thread = self._setup_autocommit_worker()
        # we need to get the most up-to-date offsets before starting consumption
        self.fetch_offsets()
        self._fetch_workers = self._setup_fetch_workers()

    def _build_default_error_handlers(self):
        def _handle_OffsetOutOfRangeError(parts):
            self._reset_offsets((owned_partition
                                 for owned_partition, pres in parts))

        return {
            UnknownTopicOrPartition.ERROR_CODE: lambda: raise_error(UnknownTopicOrPartition),
            OffsetOutOfRangeError.ERROR_CODE: _handle_OffsetOutOfRangeError,
            InvalidMessageError.ERROR_CODE: lambda: raise_error(InvalidMessageError)
        }

    @property
    def topic(self):
        return self._topic

    @property
    def partitions(self):
        return self._partitions

    def __del__(self):
        self.stop()

    def stop(self):
        self._running = False

    def _setup_autocommit_worker(self):
        def autocommitter():
            while True:
                if not self._running:
                    break
                if self._auto_commit_enable:
                    self._auto_commit()
        log.debug("Starting autocommitter thread")
        return self._cluster.handler.spawn(autocommitter)

    def _setup_fetch_workers(self):
        def fetcher():
            while True:
                if not self._running:
                    break
                self.fetch()
        return [self._cluster.handler.spawn(fetcher)
                for i in xrange(self._num_consumer_fetchers)]

    def __iter__(self):
        while True:
            message = self.consume()
            if not message and self._consumer_timed_out():
                raise StopIteration
            yield message

    def _consumer_timed_out(self):
        if self._consumer_timeout_ms == -1:
            return False
        disp = (time.time() - self._last_message_time) * 1000.0
        return disp > self._consumer_timeout_ms

    def consume(self):
        """Get one message from the consumer.

        :param timeout: Seconds to wait before returning None
        """
        owned_partition = self.partition_cycle.next()
        message = owned_partition.consume()

        if message:
            self._last_message_time = time.time()
            return message

    def _auto_commit(self):
        if not self._auto_commit_enable or self._auto_commit_interval_ms == 0:
            return

        if (time.time() - self._last_auto_commit) * 1000.0 >= self._auto_commit_interval_ms:
            log.info("Autocommitting consumer offset for consumer group %s and topic %s",
                     self._consumer_group, self._topic.name)
            self.commit_offsets()
            self._last_auto_commit = time.time()

    def commit_offsets(self):
        """Commit offsets for this consumer's topic

        Uses the offset commit/fetch API
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to commit offsets")

        reqs = [p.build_offset_commit_request() for p in self._partitions.keys()]
        log.info("Committing offsets for %d partitions", len(reqs))
        for i in xrange(self._offsets_commit_max_retries):
            if i > 0:
                log.debug("Retrying")
            time.sleep(i * (self._offsets_channel_backoff_ms / 1000))

            response = self._offset_manager.commit_consumer_group_offsets(
                self._consumer_group, 1, 'pykafka', reqs)
            parts_by_error = handle_partition_responses(
                response,
                self._default_error_handlers,
                partitions_by_id=self._partitions_by_id)
            if len(parts_by_error) == 1 and 0 in parts_by_error:
                break
            log.error("Error committing offsets for topic %s", self._topic.name)

            # retry only the partitions that errored
            parts_by_error.pop(0)
            errored_partitions = [op for err_group in parts_by_error.iteritems() for op in err_group]
            reqs = [p.build_offset_commit_request() for p in errored_partitions]

    def fetch_offsets(self):
        """Fetch offsets for this consumer's topic

        Uses the offset commit/fetch API
        Should be called when consumer starts and after any errors
        """
        if not self._consumer_group:
            raise Exception("consumer group must be specified to fetch offsets")

        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.set_offset(pres.offset)

        log.info("Fetching offsets")

        reqs = [p.build_offset_fetch_request() for p in self._partitions.keys()]
        res = self._offset_manager.fetch_consumer_group_offsets(self._consumer_group, reqs)
        handle_partition_responses(
            res,
            self._default_error_handlers,
            success_handler=_handle_success,
            partitions_by_id=self._partitions_by_id)

    def _reset_offsets(self, errored_partitions):
        """Reset offsets after an OffsetOutOfRangeError

        Issue an OffsetRequest for each partition and set the appropriate
        returned offset in the OwnedPartition per self._auto_offset_reset

        :param errored_partitions: the partitions with out-of-range offsets
        :type errored_partitions: Iterable of OwnedPartition
        """
        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.set_offset(pres.offset[0])

        # group out-of-range partitions by leader
        by_leader = defaultdict(list)
        for p in errored_partitions:
            by_leader[p.partition.leader].append(p)

        # get valid offset ranges for each partition
        for broker, owned_partitions in by_leader.iteritems():
            reqs = [owned_partition.build_offset_request(self._auto_offset_reset)
                    for owned_partition in owned_partitions]
            response = broker.request_offset_limits(reqs)
            handle_partition_responses(
                response,
                self._default_error_handlers,
                success_handler=_handle_success,
                partitions_by_id=self._partitions_by_id)

    def fetch(self):
        """Fetch new messages for all partitions

        Create a FetchRequest for each broker and send it. Enqueue each of the
        returned messages in the approprate OwnedPartition.
        """
        def _handle_success(parts):
            for owned_partition, pres in parts:
                owned_partition.enqueue_messages(pres.messages)

        for broker, owned_partitions in self._partitions_by_leader.iteritems():
            partition_reqs = []
            for owned_partition in owned_partitions:
                # attempt to acquire lock, just pass if we can't
                if owned_partition.lock.acquire(False) and \
                        owned_partition.message_count < self._queued_max_messages:
                    fetch_req = owned_partition.build_fetch_request(
                        self._fetch_message_max_bytes)
                    partition_reqs.append((owned_partition, fetch_req))
            if partition_reqs:
                response = broker.fetch_messages(
                    [a[1] for a in partition_reqs],
                    timeout=self._fetch_wait_max_ms,
                    min_bytes=self._fetch_min_bytes
                )
                handle_partition_responses(
                    response,
                    self._default_error_handlers,
                    success_handler=_handle_success,
                    partitions_by_id=self._partitions_by_id)
            for owned_partition, _ in partition_reqs:
                owned_partition.lock.release()

class OwnedPartition(object):
    """A partition that is owned by a SimpleConsumer.

    Used to keep track of offsets and the internal message queue.
    """

    def __init__(self,
                 partition,
                 consumer_group=None):
        self.partition = partition
        self._messages = Queue()
        self.last_offset_consumed = 0
        self.next_offset = 0
        self.lock = threading.Lock()

    @property
    def message_count(self):
        return self._messages.qsize()

    def set_offset(self, last_offset_consumed):
        """Set the internal offset counters from an OffsetFetchResponse

        :param last_offset_consumed: the last committed offset for this
            partition
        :type last_offset_consumed: int
        """
        self.last_offset_consumed = last_offset_consumed
        self.next_offset = last_offset_consumed + 1

    def build_offset_request(self, auto_offset_reset):
        """Create a PartitionOffsetRequest for this partition
        """
        return PartitionOffsetRequest(
            self.partition.topic.name, self.partition.id,
            auto_offset_reset, 1)

    def build_fetch_request(self, max_bytes):
        """Create a FetchPartitionRequest for this partition
        """
        return PartitionFetchRequest(
            self.partition.topic.name, self.partition.id,
            self.next_offset, max_bytes)

    def build_offset_commit_request(self):
        """Create a PartitionOffsetCommitRequest for this partition
        """
        return PartitionOffsetCommitRequest(
            self.partition.topic.name,
            self.partition.id,
            self.last_offset_consumed,
            int(time.time()),
            'pykafka'
        )

    def build_offset_fetch_request(self):
        """Create a PartitionOffsetFetchRequest for this partition
        """
        return PartitionOffsetFetchRequest(
            self.partition.topic.name,
            self.partition.id
        )

    def consume(self):
        """Get a single message from this partition
        """
        try:
            message = self._messages.get_nowait()
            self.last_offset_consumed = message.offset
            return message
        except Empty:
            return None

    def enqueue_messages(self, messages):
        for message in messages:
            if message.offset < self.last_offset_consumed:
                continue
            self._messages.put(message)
            self.next_offset = message.offset + 1
<|endoftext|>


<commit_before>import testbase
import sqlalchemy.topological as topological
from sqlalchemy import util
from testlib import *

class DependencySortTest(PersistTest):
    def assert_sort(self, tuples, node, collection=None):
        print str(node)
        def assert_tuple(tuple, node):
            if node[1]:
                cycles = node[1]
            else:
                cycles = []
            if tuple[0] is node[0] or tuple[0] in cycles:
                tuple.pop()
                if tuple[0] is node[0] or tuple[0] in cycles:
                    return
            elif len(tuple) > 1 and tuple[1] is node[0]:
                assert False, "Tuple not in dependency tree: " + str(tuple)
            for c in node[2]:
                assert_tuple(tuple, c)

        for tuple in tuples:
            assert_tuple(list(tuple), node)

        if collection is None:
            collection = []
        items = util.Set()
        def assert_unique(node):
            for item in [i for i in node[1] or [node[0]]]:
                assert item not in items
                items.add(item)
                if item in collection:
                    collection.remove(item)
            for c in node[2]:
                assert_unique(c)
        assert_unique(node)
        assert len(collection) == 0

    def testsort(self):
        rootnode = 'root'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        subnode1 = 'subnode1'
        subnode2 = 'subnode2'
        subnode3 = 'subnode3'
        subnode4 = 'subnode4'
        subsubnode1 = 'subsubnode1'
        tuples = [
            (subnode3, subsubnode1),
            (node2, subnode1),
            (node2, subnode2),
            (rootnode, node2),
            (rootnode, node3),
            (rootnode, node4),
            (node4, subnode3),
            (node4, subnode4)
        ]
        head = topological.sort_as_tree(tuples, [])
        self.assert_sort(tuples, head)

    def testsort2(self):
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        node5 = 'node5'
        node6 = 'node6'
        node7 = 'node7'
        tuples = [
            (node1, node2),
            (node3, node4),
            (node4, node5),
            (node5, node6),
            (node6, node2)
        ]
        head = topological.sort_as_tree(tuples, [node7])
        self.assert_sort(tuples, head, [node7])

    def testsort3(self):
        ['Mapper|Keyword|keywords,Mapper|IKAssociation|itemkeywords', 'Mapper|Item|items,Mapper|IKAssociation|itemkeywords']
        node1 = 'keywords'
        node2 = 'itemkeyowrds'
        node3 = 'items'
        tuples = [
            (node1, node2),
            (node3, node2),
            (node1,node3)
        ]
        head1 = topological.sort_as_tree(tuples, [node1, node2, node3])
        head2 = topological.sort_as_tree(tuples, [node3, node1, node2])
        head3 = topological.sort_as_tree(tuples, [node3, node2, node1])

        # TODO: figure out a "node == node2" function
        #self.assert_(str(head1) == str(head2) == str(head3))
        print "\n" + str(head1)
        print "\n" + str(head2)
        print "\n" + str(head3)

    def testsort4(self):
        node1 = 'keywords'
        node2 = 'itemkeyowrds'
        node3 = 'items'
        node4 = 'hoho'
        tuples = [
            (node1, node2),
            (node4, node1),
            (node1, node3),
            (node3, node2)
        ]
        head = topological.sort_as_tree(tuples, [])
        self.assert_sort(tuples, head)

    def testsort5(self):
        # this one, depenending on the weather,
        node1 = 'node1' #'00B94190'
        node2 = 'node2' #'00B94990'
        node3 = 'node3' #'00B9A9B0'
        node4 = 'node4' #'00B4F210'
        tuples = [
            (node4, node1),
            (node1, node2),
            (node4, node3),
            (node2, node3),
            (node4, node2),
            (node3, node3)
        ]
        allitems = [
            node1,
            node2,
            node3,
            node4
        ]
        head = topological.sort_as_tree(tuples, allitems, with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular(self):
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        node5 = 'node5'
        tuples = [
            (node4, node5),
            (node5, node4),
            (node1, node2),
            (node2, node3),
            (node3, node1),
            (node4, node1)
        ]
        head = topological.sort_as_tree(tuples, [], with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular2(self):
        # this condition was arising from ticket:362
        # and was not treated properly by topological sort
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        tuples = [
            (node1, node2),
            (node3, node1),
            (node2, node4),
            (node3, node2),
            (node2, node3)
        ]
        head = topological.sort_as_tree(tuples, [], with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular3(self):
        nodes = {}
        tuples = [('Question', 'Issue'), ('ProviderService', 'Issue'), ('Provider', 'Question'), ('Question', 'Provider'), ('ProviderService', 'Question'), ('Provider', 'ProviderService'), ('Question', 'Answer'), ('Issue', 'Question')]
        head = topological.sort_as_tree(tuples, [], with_cycles=True)
        self.assert_sort(tuples, head)

    def testbigsort(self):
        tuples = []
        for i in range(0,1500, 2):
            tuples.append((i, i+1))
        head = topological.sort_as_tree(tuples, [])

if __name__ == "__main__":
    testbase.main()
<commit_msg>fixed test which didnt pass along 'allitems' collection to the sort...<commit_after>import testbase
import sqlalchemy.topological as topological
from sqlalchemy import util
from testlib import *

class DependencySortTest(PersistTest):
    def assert_sort(self, tuples, node, collection=None):
        print str(node)
        def assert_tuple(tuple, node):
            if node[1]:
                cycles = node[1]
            else:
                cycles = []
            if tuple[0] is node[0] or tuple[0] in cycles:
                tuple.pop()
                if tuple[0] is node[0] or tuple[0] in cycles:
                    return
            elif len(tuple) > 1 and tuple[1] is node[0]:
                assert False, "Tuple not in dependency tree: " + str(tuple) + " " + str(node)
            for c in node[2]:
                assert_tuple(tuple, c)

        for tuple in tuples:
            assert_tuple(list(tuple), node)

        if collection is None:
            collection = []
        items = util.Set()
        def assert_unique(node):
            for item in [i for i in node[1] or [node[0]]]:
                assert item not in items
                items.add(item)
                if item in collection:
                    collection.remove(item)
            for c in node[2]:
                assert_unique(c)
        assert_unique(node)
        assert len(collection) == 0

    def testsort(self):
        rootnode = 'root'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        subnode1 = 'subnode1'
        subnode2 = 'subnode2'
        subnode3 = 'subnode3'
        subnode4 = 'subnode4'
        subsubnode1 = 'subsubnode1'
        tuples = [
            (subnode3, subsubnode1),
            (node2, subnode1),
            (node2, subnode2),
            (rootnode, node2),
            (rootnode, node3),
            (rootnode, node4),
            (node4, subnode3),
            (node4, subnode4)
        ]
        head = topological.sort_as_tree(tuples, [])
        self.assert_sort(tuples, head)

    def testsort2(self):
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        node5 = 'node5'
        node6 = 'node6'
        node7 = 'node7'
        tuples = [
            (node1, node2),
            (node3, node4),
            (node4, node5),
            (node5, node6),
            (node6, node2)
        ]
        head = topological.sort_as_tree(tuples, [node7])
        self.assert_sort(tuples, head, [node7])

    def testsort3(self):
        ['Mapper|Keyword|keywords,Mapper|IKAssociation|itemkeywords', 'Mapper|Item|items,Mapper|IKAssociation|itemkeywords']
        node1 = 'keywords'
        node2 = 'itemkeyowrds'
        node3 = 'items'
        tuples = [
            (node1, node2),
            (node3, node2),
            (node1,node3)
        ]
        head1 = topological.sort_as_tree(tuples, [node1, node2, node3])
        head2 = topological.sort_as_tree(tuples, [node3, node1, node2])
        head3 = topological.sort_as_tree(tuples, [node3, node2, node1])

        # TODO: figure out a "node == node2" function
        #self.assert_(str(head1) == str(head2) == str(head3))
        print "\n" + str(head1)
        print "\n" + str(head2)
        print "\n" + str(head3)

    def testsort4(self):
        node1 = 'keywords'
        node2 = 'itemkeyowrds'
        node3 = 'items'
        node4 = 'hoho'
        tuples = [
            (node1, node2),
            (node4, node1),
            (node1, node3),
            (node3, node2)
        ]
        head = topological.sort_as_tree(tuples, [])
        self.assert_sort(tuples, head)

    def testsort5(self):
        # this one, depenending on the weather,
        node1 = 'node1' #'00B94190'
        node2 = 'node2' #'00B94990'
        node3 = 'node3' #'00B9A9B0'
        node4 = 'node4' #'00B4F210'
        tuples = [
            (node4, node1),
            (node1, node2),
            (node4, node3),
            (node2, node3),
            (node4, node2),
            (node3, node3)
        ]
        allitems = [
            node1,
            node2,
            node3,
            node4
        ]
        head = topological.sort_as_tree(tuples, allitems, with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular(self):
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        node5 = 'node5'
        tuples = [
            (node4, node5),
            (node5, node4),
            (node1, node2),
            (node2, node3),
            (node3, node1),
            (node4, node1)
        ]
        allitems = [node1, node2, node3, node4]
        head = topological.sort_as_tree(tuples, allitems, with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular2(self):
        # this condition was arising from ticket:362
        # and was not treated properly by topological sort
        node1 = 'node1'
        node2 = 'node2'
        node3 = 'node3'
        node4 = 'node4'
        tuples = [
            (node1, node2),
            (node3, node1),
            (node2, node4),
            (node3, node2),
            (node2, node3)
        ]
        head = topological.sort_as_tree(tuples, [], with_cycles=True)
        self.assert_sort(tuples, head)

    def testcircular3(self):
        nodes = {}
        tuples = [('Question', 'Issue'), ('ProviderService', 'Issue'), ('Provider', 'Question'), ('Question', 'Provider'), ('ProviderService', 'Question'), ('Provider', 'ProviderService'), ('Question', 'Answer'), ('Issue', 'Question')]
        head = topological.sort_as_tree(tuples, [], with_cycles=True)
        self.assert_sort(tuples, head)

    def testbigsort(self):
        tuples = []
        for i in range(0,1500, 2):
            tuples.append((i, i+1))
        head = topological.sort_as_tree(tuples, [])

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


<commit_before>"""
This file collects tests relating to the many special cases we deal with when
trying to process biomolecular structures from the pdb.

These currently just test the specific cases that we've implemented fixes for
"""
import pytest
import moldesign as mdt
from .helpers import get_data_path

import parmed

from distutils.version import LooseVersion

@pytest.fixture
def pdb_3ac2():
    return mdt.read(get_data_path('3ac2.pdb'))

@pytest.fixture
def pdb_3ac2_roundtrip(pdb_3ac2):
    return mdt.read(pdb_3ac2.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_3ac2 pdb_3ac2_roundtrip'.split())
def test_missing_terminal_atoms_3ac2(request, mol):
    """ Tests that we can still detect terminal residues even if the peptide-bonded atoms
    are missing from the structure
    """
    mol = request.getfuncargvalue(mol)
    assert mol.chains['A'].n_terminal is not None
    assert mol.chains['A'].c_terminal is not None

@pytest.fixture
def pdb_1hpk():
    return mdt.read(get_data_path('1hpk.pdb'))

@pytest.fixture
def pdb_1hpk_roundtrip(pdb_1hpk):
    return mdt.read(pdb_1hpk.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_1hpk pdb_1hpk_roundtrip'.split())
def test_1hpk(request, mol):
    mol = request.getfuncargvalue(mol)
    mol = mdt.interfaces.ambertools._prep_for_tleap(mol)
    for residx in (0, 21, 49, 61, 73, 78):
        residue = mol.residues[residx]
        assert residue.resname == 'CYX'
        assert residue.name[:3] == 'CYX'

@pytest.fixture
def pdb_2jaj():
    return mdt.read(get_data_path('2jaj.pdb'))

@pytest.fixture
def pdb_2jaj_roundtrip(pdb_2jaj):
    return mdt.read(pdb_2jaj.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_2jaj pdb_2jaj_roundtrip'.split())
def test_negative_residue_numbers_2jaj(request, mol):
    if (mol == 'pdb_2jaj_roundtrip' and
            LooseVersion(getattr(parmed, '__version__', '0.0.0')) <= LooseVersion('2.7.3')):
        pytest.xfail("This test requires ParmEd 2.7.4 (not yet released as of this writing)")

    mol = request.getfuncargvalue(mol)
    res = mol.chains['B'].residues[0]
    assert res.pdbindex == -4
    assert res.index == 272
    assert res.name == 'GLY-4'

@pytest.mark.parametrize('mol', 'pdb_2jaj pdb_2jaj_roundtrip'.split())
def test_missing_residues_xtal_2jaj(request, mol):
    if mol == 'pdb_2jaj_roundtrip':
        pytest.xfail('Writing missing residue records is not yet supported.')
    mol = request.getfuncargvalue(mol)
    missingres = mol.metadata.missing_residues
    for expected in MISSINGRES_2JAJ:
        assert missingres[expected[0]][expected[2]] == expected[1]

@pytest.fixture
def pdb_1pyn():
    return mdt.read(get_data_path('1pyn.pdb'))

@pytest.fixture
def pdb_1pyn_roundtrip(pdb_1pyn):
    return mdt.read(pdb_1pyn.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_1pyn pdb_1pyn_roundtrip'.split())
def test_numeric_residue_name_1PYN(request, mol):
    """ The ligand in this residue is named "941", which causes a little trickiness
    """
    import parmed

    mol = request.getfuncargvalue(mol)
    ligand = mdt.Molecule(mol.residues[283])

    params = mdt.create_ff_parameters(ligand, charges='gasteiger')
    params.lib.put('/tmp/tmp.lib')

    contents = parmed.load_file('/tmp/tmp.lib')
    assert len(contents) == 1
    assert contents.keys()[0] == '941'

MISSINGRES_2JAJ = [('A', 'GLY', -4), ('A', 'PRO', -3), ('A', 'LEU', -2), ('A', 'GLY', -1),
                   ('A', 'MET', 0), ('A', 'ALA', 1), ('A', 'GLY', 2), ('A', 'LEU', 3),
                   ('A', 'GLY', 4), ('A', 'HIS', 5), ('A', 'PRO', 6), ('A', 'ALA', 7),
                   ('A', 'ALA', 32), ('A', 'LYS', 33), ('A', 'VAL', 282), ('A', 'ASP', 283),
                   ('A', 'SER', 284), ('B', 'GLY', 34), ('B', 'GLU', 35), ('B', 'ALA', 168),
                   ('B', 'ASP', 169), ('B', 'GLY', 170), ('B', 'VAL', 282), ('B', 'ASP', 283),
                   ('B', 'SER', 284)]

<commit_msg>Update a fragile test<commit_after>"""
This file collects tests relating to the many special cases we deal with when
trying to process biomolecular structures from the pdb.

These currently just test the specific cases that we've implemented fixes for
"""
import pytest
import moldesign as mdt
from .helpers import get_data_path

import parmed

from distutils.version import LooseVersion

@pytest.fixture
def pdb_3ac2():
    return mdt.read(get_data_path('3ac2.pdb'))

@pytest.fixture
def pdb_3ac2_roundtrip(pdb_3ac2):
    return mdt.read(pdb_3ac2.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_3ac2 pdb_3ac2_roundtrip'.split())
def test_missing_terminal_atoms_3ac2(request, mol):
    """ Tests that we can still detect terminal residues even if the peptide-bonded atoms
    are missing from the structure
    """
    mol = request.getfuncargvalue(mol)
    assert mol.chains['A'].n_terminal is not None
    assert mol.chains['A'].c_terminal is not None

@pytest.fixture
def pdb_1hpk():
    return mdt.read(get_data_path('1hpk.pdb'))

@pytest.fixture
def pdb_1hpk_roundtrip(pdb_1hpk):
    return mdt.read(pdb_1hpk.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_1hpk pdb_1hpk_roundtrip'.split())
def test_1hpk(request, mol):
    mol = request.getfuncargvalue(mol)
    mol = mdt.interfaces.ambertools._prep_for_tleap(mol)
    for residx in (0, 21, 49, 61, 73, 78):
        residue = mol.residues[residx]
        assert residue.resname == 'CYX'
        assert residue.name[:3] == 'CYX'

@pytest.fixture
def pdb_2jaj():
    return mdt.read(get_data_path('2jaj.pdb'))

@pytest.fixture
def pdb_2jaj_roundtrip(pdb_2jaj):
    return mdt.read(pdb_2jaj.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_2jaj pdb_2jaj_roundtrip'.split())
def test_negative_residue_numbers_2jaj(request, mol):
    if (mol == 'pdb_2jaj_roundtrip' and
            LooseVersion(getattr(parmed, '__version__', '0.0.0')) <= LooseVersion('2.7.3')):
        pytest.xfail("This test requires ParmEd 2.7.4 (not yet released as of this writing)")

    mol = request.getfuncargvalue(mol)
    res = mol.chains['B'].residues[0]
    assert res.pdbindex == -4
    assert res.index == 272
    assert res.name == 'GLY-4'

@pytest.mark.parametrize('mol', 'pdb_2jaj pdb_2jaj_roundtrip'.split())
def test_missing_residues_xtal_2jaj(request, mol):
    if mol == 'pdb_2jaj_roundtrip':
        pytest.xfail('Writing missing residue records is not yet supported.')
    mol = request.getfuncargvalue(mol)
    missingres = mol.metadata.missing_residues
    for expected in MISSINGRES_2JAJ:
        assert missingres[expected[0]][expected[2]] == expected[1]

@pytest.fixture
def pdb_1pyn():
    return mdt.read(get_data_path('1pyn.pdb'))

@pytest.fixture
def pdb_1pyn_roundtrip(pdb_1pyn):
    return mdt.read(pdb_1pyn.write('pdb'), format='pdb')

@pytest.mark.parametrize('mol', 'pdb_1pyn pdb_1pyn_roundtrip'.split())
def test_numeric_residue_name_1PYN(request, mol):
    """ The ligand in this residue is named "941", which causes a little trickiness
    """
    import parmed

    mol = request.getfuncargvalue(mol)
    ligand = mdt.Molecule(mol.residues[283])

    params = mdt.create_ff_parameters(ligand, charges='gasteiger')
    params._file_list['mol.lib'].put('/tmp/tmp.lib')

    contents = parmed.load_file('/tmp/tmp.lib')
    assert len(contents) == 1
    assert contents.keys()[0] == '941'

MISSINGRES_2JAJ = [('A', 'GLY', -4), ('A', 'PRO', -3), ('A', 'LEU', -2), ('A', 'GLY', -1),
                   ('A', 'MET', 0), ('A', 'ALA', 1), ('A', 'GLY', 2), ('A', 'LEU', 3),
                   ('A', 'GLY', 4), ('A', 'HIS', 5), ('A', 'PRO', 6), ('A', 'ALA', 7),
                   ('A', 'ALA', 32), ('A', 'LYS', 33), ('A', 'VAL', 282), ('A', 'ASP', 283),
                   ('A', 'SER', 284), ('B', 'GLY', 34), ('B', 'GLU', 35), ('B', 'ALA', 168),
                   ('B', 'ASP', 169), ('B', 'GLY', 170), ('B', 'VAL', 282), ('B', 'ASP', 283),
                   ('B', 'SER', 284)]
<|endoftext|>