<s> from __future__ import absolute_import <EOL> import os <EOL> import functools <EOL> import itertools <EOL> import fnmatch <EOL> from collections import namedtuple <EOL> from opencorpora import compat , xml_utils <EOL> from opencorpora . compat import imap , text_type <EOL> def make_iterable ( obj , default = None ) : <EOL> """<STR_LIT>""" <EOL> if obj is None : <EOL> return default or [ ] <EOL> if isinstance ( obj , ( compat . string_types , compat . integer_types ) ) : <EOL> return [ obj ] <EOL> return obj <EOL> def some_items_match ( items , patterns ) : <EOL> return any ( <EOL> fnmatch . fnmatchcase ( item , pattern ) <EOL> for item in items <EOL> for pattern in patterns <EOL> ) <EOL> def _sentence_source ( sent_elem ) : <EOL> return text_type ( sent_elem . find ( '<STR_LIT:source>' ) . text ) <EOL> def _sentence_words ( sent_elem ) : <EOL> return [ text_type ( tok . get ( '<STR_LIT:text>' ) ) for tok in sent_elem . findall ( '<STR_LIT>' ) ] <EOL> def _sentence_tagged_words ( sent_elem ) : <EOL> res = [ ] <EOL> for tok in sent_elem . findall ( '<STR_LIT>' ) : <EOL> text = text_type ( tok . get ( '<STR_LIT:text>' ) ) <EOL> parse = tok . find ( '<STR_LIT>' ) <EOL> tag = text_type ( '<STR_LIT:U+002C>' ) . join ( _grammemes ( parse ) ) <EOL> res . append ( ( text , tag ) ) <EOL> return res <EOL> def _sentence_parsed_words ( sent_elem ) : <EOL> res = [ ] <EOL> for tok in sent_elem . findall ( '<STR_LIT>' ) : <EOL> text = text_type ( tok . get ( '<STR_LIT:text>' ) ) <EOL> parses = tok . findall ( '<STR_LIT>' ) <EOL> annotations = [ <EOL> ( text_type ( p . get ( '<STR_LIT:t>' ) ) , text_type ( '<STR_LIT:U+002C>' ) . join ( _grammemes ( p ) ) ) <EOL> for p in parses <EOL> ] <EOL> res . append ( ( text , annotations ) ) <EOL> return res <EOL> def _grammemes ( l_element ) : <EOL> return [ text_type ( grammeme . get ( '<STR_LIT:v>' ) ) <EOL> for grammeme in l_element . getchildren ( ) ] <EOL> def non_iterative ( func ) : <EOL> @ functools . wraps ( func ) <EOL> def res ( * args , ** kwargs ) : <EOL> return list ( func ( * args , ** kwargs ) ) <EOL> return res <EOL> class Document ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , xml ) : <EOL> self . root = xml <EOL> def _xml_sents ( self ) : <EOL> return self . root . findall ( '<STR_LIT>' ) <EOL> def _xml_paras ( self ) : <EOL> return self . root . findall ( '<STR_LIT>' ) <EOL> def iter_sents ( self ) : <EOL> return imap ( _sentence_words , self . _xml_sents ( ) ) <EOL> def iter_raw_sents ( self ) : <EOL> return imap ( _sentence_source , self . _xml_sents ( ) ) <EOL> def iter_tagged_sents ( self ) : <EOL> return imap ( _sentence_tagged_words , self . _xml_sents ( ) ) <EOL> def iter_parsed_sents ( self ) : <EOL> return imap ( _sentence_parsed_words , self . _xml_sents ( ) ) <EOL> def iter_paras ( self ) : <EOL> for para_elem in self . _xml_paras ( ) : <EOL> yield [ _sentence_words ( s ) for s in para_elem . findall ( '<STR_LIT>' ) ] <EOL> def iter_raw_paras ( self ) : <EOL> for para_elem in self . _xml_paras ( ) : <EOL> yield "<STR_LIT:U+0020>" . join ( _sentence_source ( s ) for s in para_elem . findall ( '<STR_LIT>' ) ) <EOL> def iter_tagged_paras ( self ) : <EOL> for para_elem in self . _xml_paras ( ) : <EOL> yield [ _sentence_tagged_words ( s ) for s in para_elem . findall ( '<STR_LIT>' ) ] <EOL> def iter_parsed_paras ( self ) : <EOL> for para_elem in self . _xml_paras ( ) : <EOL> yield [ _sentence_parsed_words ( s ) for s in para_elem . findall ( '<STR_LIT>' ) ] <EOL> def iter_words ( self ) : <EOL> return itertools . chain ( * self . iter_sents ( ) ) <EOL> def iter_tagged_words ( self ) : <EOL> return itertools . chain ( * self . iter_tagged_sents ( ) ) <EOL> def iter_parsed_words ( self ) : <EOL> return itertools . chain ( * self . iter_parsed_sents ( ) ) <EOL> sents = non_iterative ( iter_sents ) <EOL> raw_sents = non_iterative ( iter_raw_sents ) <EOL> tagged_sents = non_iterative ( iter_tagged_sents ) <EOL> parsed_sents = non_iterative ( iter_parsed_sents ) <EOL> paras = non_iterative ( iter_paras ) <EOL> raw_paras = non_iterative ( iter_raw_paras ) <EOL> tagged_paras = non_iterative ( iter_tagged_paras ) <EOL> parsed_paras = non_iterative ( iter_parsed_paras ) <EOL> words = non_iterative ( iter_words ) <EOL> tagged_words = non_iterative ( iter_tagged_words ) <EOL> parsed_words = non_iterative ( iter_parsed_words ) <EOL> def title ( self ) : <EOL> return text_type ( self . root . get ( '<STR_LIT:name>' ) ) <EOL> def categories ( self ) : <EOL> return [ text_type ( tag . text ) for tag in self . root . findall ( '<STR_LIT>' ) ] <EOL> def raw ( self ) : <EOL> return "<STR_LIT>" . join ( self . iter_raw_paras ( ) ) <EOL> @ compat . utf8_for_PY2 <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , self . title ( ) ) <EOL> def destroy ( self ) : <EOL> self . root . clear ( ) <EOL> _DocumentMeta = namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> def _from_documents ( doc_method_name ) : <EOL> def method ( self , fileids = None , categories = None ) : <EOL> return self . _doc_iterator ( fileids , categories , doc_method_name ) <EOL> method . __name__ = str ( doc_method_name ) <EOL> return method <EOL> class CorpusReader ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , filename , cache_filename = None , use_cache = True ) : <EOL> self . filename = filename <EOL> self . use_cache = use_cache <EOL> self . _document_meta = None <EOL> self . _cache_filename = cache_filename or filename + '<STR_LIT>' <EOL> def iter_documents ( self , fileids = None , categories = None , _destroy = False ) : <EOL> """<STR_LIT>""" <EOL> doc_ids = self . _filter_ids ( fileids , categories ) <EOL> for doc in imap ( self . get_document , doc_ids ) : <EOL> yield doc <EOL> if _destroy : <EOL> doc . destroy ( ) <EOL> def iter_documents_raw ( self , fileids = None , categories = None ) : <EOL> for doc in self . iter_documents ( fileids , categories , _destroy = True ) : <EOL> yield doc . raw ( ) <EOL> iter_paras = _from_documents ( '<STR_LIT>' ) <EOL> iter_raw_paras = _from_documents ( '<STR_LIT>' ) <EOL> iter_tagged_paras = _from_documents ( '<STR_LIT>' ) <EOL> iter_parsed_paras = _from_documents ( '<STR_LIT>' ) <EOL> iter_sents = _from_documents ( '<STR_LIT>' ) <EOL> iter_raw_sents = _from_documents ( '<STR_LIT>' ) <EOL> iter_tagged_sents = _from_documents ( '<STR_LIT>' ) <EOL> iter_parsed_sents = _from_documents ( '<STR_LIT>' ) <EOL> iter_words = _from_documents ( '<STR_LIT>' ) <EOL> iter_tagged_words = _from_documents ( '<STR_LIT>' ) <EOL> iter_parsed_words = _from_documents ( '<STR_LIT>' ) <EOL> sents = non_iterative ( iter_sents ) <EOL> raw_sents = non_iterative ( iter_raw_sents ) <EOL> tagged_sents = non_iterative ( iter_tagged_sents ) <EOL> parsed_sents = non_iterative ( iter_parsed_sents ) <EOL> paras = non_iterative ( iter_paras ) <EOL> raw_paras = non_iterative ( iter_raw_paras ) <EOL> tagged_paras = non_iterative ( iter_tagged_paras ) <EOL> parsed_paras = non_iterative ( iter_parsed_paras ) <EOL> words = non_iterative ( iter_words ) <EOL> tagged_words = non_iterative ( iter_tagged_words ) <EOL> parsed_words = non_iterative ( iter_parsed_words ) <EOL> documents = non_iterative ( iter_documents ) <EOL> documents_raw = non_iterative ( iter_documents_raw ) <EOL> def raw ( self , fileids = None , categories = None ) : <EOL> return "<STR_LIT>" . join ( self . iter_documents_raw ( fileids , categories ) ) <EOL> def _doc_iterator ( self , fileids , categories , doc_method ) : <EOL> for doc in self . iter_documents ( fileids , categories , _destroy = True ) : <EOL> meth = getattr ( doc , doc_method ) <EOL> for res in meth ( ) : <EOL> yield res <EOL> def fileids ( self , categories = None ) : <EOL> return list ( self . _filter_ids ( None , categories ) ) <EOL> def categories ( self , fileids = None , patterns = None ) : <EOL> meta = self . _get_meta ( ) <EOL> fileids = make_iterable ( fileids , meta . keys ( ) ) <EOL> result = sorted ( set ( <EOL> cat for cat in itertools . chain ( * ( <EOL> meta [ str ( doc_id ) ] . categories for doc_id in fileids <EOL> ) ) <EOL> ) ) <EOL> if patterns : <EOL> patterns = make_iterable ( patterns ) <EOL> result = [ cat for cat in result <EOL> if some_items_match ( [ cat ] , patterns ) ] <EOL> return result <EOL> def catalog ( self , categories = None ) : <EOL> """<STR_LIT>""" <EOL> ids = self . _filter_ids ( None , categories ) <EOL> doc_meta = self . _get_meta ( ) <EOL> return [ ( doc_id , doc_meta [ doc_id ] . title ) for doc_id in ids ] <EOL> def get_document ( self , doc_id ) : <EOL> """<STR_LIT>""" <EOL> return Document ( self . _document_xml ( doc_id ) ) <EOL> def readme ( self ) : <EOL> return self . __doc__ <EOL> def get_annotation_info ( self ) : <EOL> for el in xml_utils . iterparse ( self . filename , '<STR_LIT>' , clear = True , <EOL> events = ( '<STR_LIT:start>' , ) ) : <EOL> return { <EOL> '<STR_LIT:version>' : el . get ( '<STR_LIT:version>' ) , <EOL> '<STR_LIT>' : el . get ( '<STR_LIT>' ) , <EOL> } <EOL> def _filter_ids ( self , fileids = None , categories = None ) : <EOL> meta = self . _get_meta ( ) <EOL> fileids = make_iterable ( fileids , meta . keys ( ) ) <EOL> if categories is None : <EOL> return imap ( str , fileids ) <EOL> category_patterns = make_iterable ( categories ) <EOL> return ( doc_id for doc_id in fileids <EOL> if some_items_match ( <EOL> meta [ doc_id ] . categories , category_patterns ) ) <EOL> def _get_meta ( self ) : <EOL> if self . _document_meta is None and self . use_cache : <EOL> self . _load_meta_cache ( ) <EOL> if self . _document_meta is None : <EOL> self . _document_meta = self . _compute_document_meta ( ) <EOL> if self . use_cache : <EOL> self . _create_meta_cache ( ) <EOL> return self . _document_meta <EOL> def _create_meta_cache ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> with open ( self . _cache_filename , '<STR_LIT:wb>' ) as f : <EOL> compat . pickle . dump ( self . _document_meta , f , <NUM_LIT:1> ) <EOL> except ( IOError , compat . pickle . PickleError ) : <EOL> pass <EOL> def _load_meta_cache ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> if self . _should_invalidate_cache ( ) : <EOL> os . remove ( self . _cache_filename ) <EOL> else : <EOL> with open ( self . _cache_filename , '<STR_LIT:rb>' ) as f : <EOL> self . _document_meta = compat . pickle . load ( f ) <EOL> except ( OSError , IOError , compat . pickle . PickleError , <EOL> ImportError , AttributeError ) : <EOL> pass <EOL> def _should_invalidate_cache ( self ) : <EOL> data_mtime = os . path . getmtime ( self . filename ) <EOL> cache_mtime = os . path . getmtime ( self . _cache_filename ) <EOL> return data_mtime > cache_mtime <EOL> def _compute_document_meta ( self ) : <EOL> """<STR_LIT>""" <EOL> meta = compat . OrderedDict ( ) <EOL> bounds_iter = xml_utils . bounds ( self . filename , <EOL> start_re = r'<STR_LIT>' , <EOL> end_re = r'<STR_LIT>' ) <EOL> for match , bounds in bounds_iter : <EOL> doc_id , title = str ( match . group ( <NUM_LIT:1> ) ) , match . group ( <NUM_LIT:2> ) <EOL> title = xml_utils . unescape_attribute ( title ) <EOL> xml_data = xml_utils . load_chunk ( self . filename , bounds ) <EOL> doc = Document ( compat . ElementTree . XML ( xml_data . encode ( '<STR_LIT:utf8>' ) ) ) <EOL> meta [ doc_id ] = _DocumentMeta ( title , bounds , doc . categories ( ) ) <EOL> return meta <EOL> def _document_xml ( self , doc_id ) : <EOL> """<STR_LIT>""" <EOL> doc_str = self . _get_doc_by_raw_offset ( str ( doc_id ) ) <EOL> return compat . ElementTree . XML ( doc_str . encode ( '<STR_LIT:utf8>' ) ) <EOL> def _get_doc_by_raw_offset ( self , doc_id ) : <EOL> """<STR_LIT>""" <EOL> bounds = self . _get_meta ( ) [ str ( doc_id ) ] . bounds <EOL> return xml_utils . load_chunk ( self . filename , bounds ) <EOL> def _get_doc_by_line_offset ( self , doc_id ) : <EOL> """<STR_LIT>""" <EOL> bounds = self . _get_meta ( ) [ str ( doc_id ) ] . bounds <EOL> return xml_utils . load_chunk ( self . filename , bounds , slow = True ) </s>
