<commit_before>##
# Author: Ajay Thampi
##
from __future__ import print_function
import os
import sys
import csv
if sys.platform == 'win32':
    # Windows C long is 32 bits, and the Python int is too large to fit inside.
    # Use the limit appropriate for a 32-bit integer as the max file size
    csv.field_size_limit(2**31-1)
else:
    csv.field_size_limit(sys.maxsize)
import zipfile
from scipy.spatial import cKDTree as KDTree
from reverse_geocoder import cKDTree_MP as KDTree_MP
import numpy as np

GN_URL = 'http://download.geonames.org/export/dump/'
GN_CITIES1000 = 'cities1000'
GN_ADMIN1 = 'admin1CodesASCII.txt'
GN_ADMIN2 = 'admin2Codes.txt'

GN_COLUMNS = {
    'geoNameId': 0,
    'name': 1,
    'asciiName': 2,
    'alternateNames': 3,
    'latitude': 4,
    'longitude': 5,
    'featureClass': 6,
    'featureCode': 7,
    'countryCode': 8,
    'cc2': 9,
    'admin1Code': 10,
    'admin2Code': 11,
    'admin3Code': 12,
    'admin4Code': 13,
    'population': 14,
    'elevation': 15,
    'dem': 16,
    'timezone': 17,
    'modificationDate': 18
}

ADMIN_COLUMNS = {
    'concatCodes': 0,
    'name': 1,
    'asciiName': 2,
    'geoNameId': 3
}

RG_COLUMNS = [
    'lat',
    'lon',
    'name',
    'admin1',
    'admin2',
    'cc'
]

RG_FILE = 'rg_cities1000.csv'

A = 6378.137 # major axis in kms
E2 = 0.00669437999014

def singleton(cls):
    instances = {}
    def getinstance(mode=2,verbose=True,stream=None):
        if cls not in instances:
            instances[cls] = cls(mode=mode,verbose=verbose,stream=stream)
        return instances[cls]
    return getinstance

@singleton
class RGeocoder:
    def __init__(self,mode=2,verbose=True,stream=None):
        self.mode = mode
        self.verbose = verbose
        if stream:
            coordinates, self.locations = self.load(stream)
        else:
            coordinates, self.locations = self.extract(rel_path(RG_FILE))

        if mode == 1: # Single-process
            self.tree = KDTree(coordinates)
        else: # Multi-process
            self.tree = KDTree_MP.cKDTree_MP(coordinates)

    def query(self,coordinates):
        try:
            if self.mode == 1:
                distances,indices = self.tree.query(coordinates,k=1)
            else:
                distances,indices = self.tree.pquery(coordinates,k=1)
        except ValueError as e:
            raise e
        else:
            return [self.locations[index] for index in indices]

    def load(self, stream):
        stream_reader = csv.DictReader(stream, delimiter=',')
        header = stream_reader.fieldnames

        if header != RG_COLUMNS:
            raise csv.Error('Input must be a comma-separated file with header containing ' + \
                'the following columns - %s. For more help, visit: ' % (','.join(RG_COLUMNS)) + \
                'https://github.com/thampiman/reverse-geocoder' )

        # Load all the coordinates and locations
        geo_coords,locations = [],[]
        for row in stream_reader:
            geo_coords.append((row['lat'],row['lon']))
            locations.append(row)

        return geo_coords,locations

    def extract(self,local_filename):
        if os.path.exists(local_filename):
            if self.verbose:
                print('Loading formatted geocoded file...')
            rows = csv.DictReader(open(local_filename,'rt'))
        else:
            gn_cities1000_url = GN_URL + GN_CITIES1000 + '.zip'
            gn_admin1_url = GN_URL + GN_ADMIN1
            gn_admin2_url = GN_URL + GN_ADMIN2

            cities1000_zipfilename = GN_CITIES1000 + '.zip'
            cities1000_filename = GN_CITIES1000 + '.txt'

            if not os.path.exists(cities1000_zipfilename):
                if self.verbose:
                    print('Downloading files from Geoname...')
                try: # Python 3
                    import urllib.request
                    urllib.request.urlretrieve(gn_cities1000_url,cities1000_zipfilename)
                    urllib.request.urlretrieve(gn_admin1_url,GN_ADMIN1)
                    urllib.request.urlretrieve(gn_admin2_url,GN_ADMIN2)
                except ImportError: # Python 2
                    import urllib
                    urllib.urlretrieve(gn_cities1000_url,cities1000_zipfilename)
                    urllib.urlretrieve(gn_admin1_url,GN_ADMIN1)
                    urllib.urlretrieve(gn_admin2_url,GN_ADMIN2)

            if self.verbose:
                print('Extracting cities1000...')
            z = zipfile.ZipFile(open(cities1000_zipfilename,'rb'))
            open(cities1000_filename,'wb').write(z.read(cities1000_filename))

            if self.verbose:
                print('Loading admin1 codes...')
            admin1_map = {}
            t_rows = csv.reader(open(GN_ADMIN1,'rt'),delimiter='\t')
            for row in t_rows:
                admin1_map[row[ADMIN_COLUMNS['concatCodes']]] = row[ADMIN_COLUMNS['asciiName']]

            if self.verbose:
                print('Loading admin2 codes...')
            admin2_map = {}
            for row in csv.reader(open(GN_ADMIN2,'rt'),delimiter='\t'):
                admin2_map[row[ADMIN_COLUMNS['concatCodes']]] = row[ADMIN_COLUMNS['asciiName']]

            if self.verbose:
                print('Creating formatted geocoded file...')
            writer = csv.DictWriter(open(local_filename,'wt'),fieldnames=RG_COLUMNS)
            rows = []
            for row in csv.reader(open(cities1000_filename,'rt'),delimiter='\t',quoting=csv.QUOTE_NONE):
                lat = row[GN_COLUMNS['latitude']]
                lon = row[GN_COLUMNS['longitude']]
                name = row[GN_COLUMNS['asciiName']]
                cc = row[GN_COLUMNS['countryCode']]

                admin1_c = row[GN_COLUMNS['admin1Code']]
                admin2_c = row[GN_COLUMNS['admin2Code']]

                cc_admin1 = cc+'.'+admin1_c
                cc_admin2 = cc+'.'+admin1_c+'.'+admin2_c

                admin1 = ''
                admin2 = ''

                if cc_admin1 in admin1_map:
                    admin1 = admin1_map[cc_admin1]
                if cc_admin2 in admin2_map:
                    admin2 = admin2_map[cc_admin2]

                write_row = {'lat':lat,'lon':lon,'name':name,'admin1':admin1,'admin2':admin2,'cc':cc}
                rows.append(write_row)
            writer.writeheader()
            writer.writerows(rows)

            if self.verbose:
                print('Removing extracted cities1000 to save space...')
            os.remove(cities1000_filename)

        # Load all the coordinates and locations
        geo_coords,locations = [],[]
        for row in rows:
            geo_coords.append((row['lat'],row['lon']))
            locations.append(row)
        ecef_coords = geodetic_in_ecef(geo_coords)
        return geo_coords,locations

def geodetic_in_ecef(geo_coords):
    geo_coords = np.asarray(geo_coords).astype(np.float)
    lat = geo_coords[:,0]
    lon = geo_coords[:,1]

    lat_r = np.radians(lat)
    lon_r = np.radians(lon)
    normal = A / (np.sqrt(1 - E2*(np.sin(lat_r) ** 2)))

    x = normal * np.cos(lat_r) * np.cos(lon_r)
    y = normal * np.cos(lat_r) * np.sin(lon_r)
    z = normal * (1 - E2) * np.sin(lat)

    return np.column_stack([x,y,z])

def rel_path(filename):
    return os.path.join(os.getcwd(), os.path.dirname(__file__), filename)

def get(geo_coord,mode=2,verbose=True):
    if type(geo_coord) != tuple or type(geo_coord[0]) != float:
        raise TypeError('Expecting a tuple')

    rg = RGeocoder(mode=mode,verbose=verbose)
    return rg.query([geo_coord])[0]

def search(geo_coords,mode=2,verbose=True):
    if type(geo_coords) != tuple and type(geo_coords) != list:
        raise TypeError('Expecting a tuple or a tuple/list of tuples')
    elif type(geo_coords[0]) != tuple:
        geo_coords = [geo_coords]

    rg = RGeocoder(mode=mode,verbose=verbose)
    return rg.query(geo_coords)

if __name__ == '__main__':
    print('Testing single coordinate through get...')
    city = (37.78674,-122.39222)
    print('Reverse geocoding 1 city...')
    result = get(city)
    print(result)

    print('Testing coordinates...')
    cities = [(51.5214588,-0.1729636),(9.936033, 76.259952),(37.38605,-122.08385)]
    print('Reverse geocoding %d cities...' % len(cities))
    results = search(cities)
    print(results)
<commit_msg>Pylint fixes + docstrings<commit_after>""" A Fast, Offline Reverse Geocoder in Python

A Python library for offline reverse geocoding. It improves on an existing library
called reverse_geocode developed by Richard Penman.
"""
from __future__ import print_function

__author__ = 'Ajay Thampi'
import os
import sys
import csv
if sys.platform == 'win32':
    # Windows C long is 32 bits, and the Python int is too large to fit inside.
    # Use the limit appropriate for a 32-bit integer as the max file size
    csv.field_size_limit(2**31-1)
else:
    csv.field_size_limit(sys.maxsize)
import zipfile
from scipy.spatial import cKDTree as KDTree
from reverse_geocoder import cKDTree_MP as KDTree_MP
import numpy as np

GN_URL = 'http://download.geonames.org/export/dump/'
GN_CITIES1000 = 'cities1000'
GN_ADMIN1 = 'admin1CodesASCII.txt'
GN_ADMIN2 = 'admin2Codes.txt'

# Schema of the GeoNames Cities with Population > 1000
GN_COLUMNS = {
    'geoNameId': 0,
    'name': 1,
    'asciiName': 2,
    'alternateNames': 3,
    'latitude': 4,
    'longitude': 5,
    'featureClass': 6,
    'featureCode': 7,
    'countryCode': 8,
    'cc2': 9,
    'admin1Code': 10,
    'admin2Code': 11,
    'admin3Code': 12,
    'admin4Code': 13,
    'population': 14,
    'elevation': 15,
    'dem': 16,
    'timezone': 17,
    'modificationDate': 18
}

# Schema of the GeoNames Admin 1/2 Codes
ADMIN_COLUMNS = {
    'concatCodes': 0,
    'name': 1,
    'asciiName': 2,
    'geoNameId': 3
}

# Schema of the cities file created by this library
RG_COLUMNS = [
    'lat',
    'lon',
    'name',
    'admin1',
    'admin2',
    'cc'
]

# Name of cities file created by this library
RG_FILE = 'rg_cities1000.csv'

# WGS-84 major axis in kms
A = 6378.137

# WGS-84 eccentricity squared
E2 = 0.00669437999014

def singleton(cls):
    """
    Function to get single instance of the RGeocoder class
    """
    instances = {}
    def getinstance(mode=2, verbose=True, stream=None):
        """
        Creates a new RGeocoder instance if not created already
        """
        if cls not in instances:
            instances[cls] = cls(mode=mode, verbose=verbose, stream=stream)
        return instances[cls]
    return getinstance

@singleton
class RGeocoder(object):
    """
    The main reverse geocoder class
    """
    def __init__(self, mode=2, verbose=True, stream=None):
        """ Class Instantiation
        Args:
        mode (int): Library supports the following two modes:
                    - 1 = Single-threaded K-D Tree
                    - 2 = Multi-threaded K-D Tree (Default)
        verbose (bool): For verbose output, set to True
        stream (io.StringIO): An in-memory stream of a custom data source
        """
        self.mode = mode
        self.verbose = verbose
        if stream:
            coordinates, self.locations = self.load(stream)
        else:
            coordinates, self.locations = self.extract(rel_path(RG_FILE))

        if mode == 1: # Single-process
            self.tree = KDTree(coordinates)
        else: # Multi-process
            self.tree = KDTree_MP.cKDTree_MP(coordinates)

    def query(self, coordinates):
        """
        Function to query the K-D tree to find the nearest city
        Args:
        coordinates (list): List of tuple coordinates, i.e. [(latitude, longitude)]
        """
        try:
            if self.mode == 1:
                _, indices = self.tree.query(coordinates, k=1)
            else:
                _, indices = self.tree.pquery(coordinates, k=1)
        except ValueError as _e:
            raise _e
        else:
            return [self.locations[index] for index in indices]

    def load(self, stream):
        """
        Function that loads a custom data source
        Args:
        stream (io.StringIO): An in-memory stream of a custom data source.
                              The format of the stream must be a comma-separated file
                              with header containing the columns defined in RG_COLUMNS.
        """
        stream_reader = csv.DictReader(stream, delimiter=',')
        header = stream_reader.fieldnames

        if header != RG_COLUMNS:
            raise csv.Error('Input must be a comma-separated file with header containing ' + \
                'the following columns - %s. For more help, visit: ' % (','.join(RG_COLUMNS)) + \
                'https://github.com/thampiman/reverse-geocoder')

        # Load all the coordinates and locations
        geo_coords, locations = [], []
        for row in stream_reader:
            geo_coords.append((row['lat'], row['lon']))
            locations.append(row)

        return geo_coords, locations

    def extract(self, local_filename):
        """
        Function loads the already extracted GeoNames cities file or downloads and extracts it if
        it doesn't exist locally
        Args:
        local_filename (str): Path to local RG_FILE
        """
        if os.path.exists(local_filename):
            if self.verbose:
                print('Loading formatted geocoded file...')
            rows = csv.DictReader(open(local_filename, 'rt'))
        else:
            gn_cities1000_url = GN_URL + GN_CITIES1000 + '.zip'
            gn_admin1_url = GN_URL + GN_ADMIN1
            gn_admin2_url = GN_URL + GN_ADMIN2

            cities1000_zipfilename = GN_CITIES1000 + '.zip'
            cities1000_filename = GN_CITIES1000 + '.txt'

            if not os.path.exists(cities1000_zipfilename):
                if self.verbose:
                    print('Downloading files from Geoname...')
                try: # Python 3
                    import urllib.request
                    urllib.request.urlretrieve(gn_cities1000_url, cities1000_zipfilename)
                    urllib.request.urlretrieve(gn_admin1_url, GN_ADMIN1)
                    urllib.request.urlretrieve(gn_admin2_url, GN_ADMIN2)
                except ImportError: # Python 2
                    import urllib
                    urllib.urlretrieve(gn_cities1000_url, cities1000_zipfilename)
                    urllib.urlretrieve(gn_admin1_url, GN_ADMIN1)
                    urllib.urlretrieve(gn_admin2_url, GN_ADMIN2)

            if self.verbose:
                print('Extracting cities1000...')
            _z = zipfile.ZipFile(open(cities1000_zipfilename, 'rb'))
            open(cities1000_filename, 'wb').write(_z.read(cities1000_filename))

            if self.verbose:
                print('Loading admin1 codes...')
            admin1_map = {}
            t_rows = csv.reader(open(GN_ADMIN1, 'rt'), delimiter='\t')
            for row in t_rows:
                admin1_map[row[ADMIN_COLUMNS['concatCodes']]] = row[ADMIN_COLUMNS['asciiName']]

            if self.verbose:
                print('Loading admin2 codes...')
            admin2_map = {}
            for row in csv.reader(open(GN_ADMIN2, 'rt'), delimiter='\t'):
                admin2_map[row[ADMIN_COLUMNS['concatCodes']]] = row[ADMIN_COLUMNS['asciiName']]

            if self.verbose:
                print('Creating formatted geocoded file...')
            writer = csv.DictWriter(open(local_filename, 'wt'), fieldnames=RG_COLUMNS)
            rows = []
            for row in csv.reader(open(cities1000_filename, 'rt'), \
                    delimiter='\t', quoting=csv.QUOTE_NONE):
                lat = row[GN_COLUMNS['latitude']]
                lon = row[GN_COLUMNS['longitude']]
                name = row[GN_COLUMNS['asciiName']]
                cc = row[GN_COLUMNS['countryCode']]

                admin1_c = row[GN_COLUMNS['admin1Code']]
                admin2_c = row[GN_COLUMNS['admin2Code']]

                cc_admin1 = cc+'.'+admin1_c
                cc_admin2 = cc+'.'+admin1_c+'.'+admin2_c

                admin1 = ''
                admin2 = ''

                if cc_admin1 in admin1_map:
                    admin1 = admin1_map[cc_admin1]
                if cc_admin2 in admin2_map:
                    admin2 = admin2_map[cc_admin2]

                write_row = {'lat':lat,
                             'lon':lon,
                             'name':name,
                             'admin1':admin1,
                             'admin2':admin2,
                             'cc':cc}
                rows.append(write_row)
            writer.writeheader()
            writer.writerows(rows)

            if self.verbose:
                print('Removing extracted cities1000 to save space...')
            os.remove(cities1000_filename)

        # Load all the coordinates and locations
        geo_coords, locations = [], []
        for row in rows:
            geo_coords.append((row['lat'], row['lon']))
            locations.append(row)
        return geo_coords, locations

def geodetic_in_ecef(geo_coords):
    geo_coords = np.asarray(geo_coords).astype(np.float)
    lat = geo_coords[:, 0]
    lon = geo_coords[:, 1]

    lat_r = np.radians(lat)
    lon_r = np.radians(lon)
    normal = A / (np.sqrt(1 - E2 * (np.sin(lat_r) ** 2)))

    x = normal * np.cos(lat_r) * np.cos(lon_r)
    y = normal * np.cos(lat_r) * np.sin(lon_r)
    z = normal * (1 - E2) * np.sin(lat)

    return np.column_stack([x, y, z])

def rel_path(filename):
    """
    Function that gets relative path to the filename
    """
    return os.path.join(os.getcwd(), os.path.dirname(__file__), filename)

def get(geo_coord, mode=2, verbose=True):
    """
    Function to query for a single coordinate
    """
    if not isinstance(geo_coord, tuple) or not isinstance(geo_coord[0], float):
        raise TypeError('Expecting a tuple')

    _rg = RGeocoder(mode=mode, verbose=verbose)
    return _rg.query([geo_coord])[0]

def search(geo_coords, mode=2, verbose=True):
    """
    Function to query for a list of coordinates
    """
    if not isinstance(geo_coords, tuple) and not isinstance(geo_coords, list):
        raise TypeError('Expecting a tuple or a tuple/list of tuples')
    elif not isinstance(geo_coords[0], tuple):
        geo_coords = [geo_coords]

    _rg = RGeocoder(mode=mode, verbose=verbose)
    return _rg.query(geo_coords)

if __name__ == '__main__':
    print('Testing single coordinate through get...')
    city = (37.78674, -122.39222)
    print('Reverse geocoding 1 city...')
    result = get(city)
    print(result)

    print('Testing coordinates...')
    cities = [(51.5214588, -0.1729636), (9.936033, 76.259952), (37.38605, -122.08385)]
    print('Reverse geocoding %d cities...' % len(cities))
    results = search(cities)
    print(results)
<|endoftext|>


<commit_before>from xml.dom import minidom
from types import *
import datetime
import csv
import os
import re
import sqlite3
import copy
import unicodedata

class XMLPatent(object):
    """
    Base class for handling parsing of xml fields
    """

    def __init__(self, xmlstring):
        """
        [xmlstring]: string containing xml document. Should be provided by
        parse.py
        """
        # lowercase all tags
        self.xmlstring = re.sub(r"<[/]?[A-Za-z-]+?[ >]", lambda x: x.group().lower(), xmlstring)
        # store the minidom parsed xml doc
        self.xmldoc = minidom.parseString(self.xmlstring)

        # country, patent, kind, date_grant
        self.country, self.patent, self.kind, self.date_grant = self.__tagNme__(xmldoc, ["publication-reference", ["country", "doc-number", "kind", "date"]])
        # pat_type
        self.pat_type = self.__tagNme__(xmldoc, ["application-reference"], iHTML=False)[0].attributes["appl-type"].value
        # date_app, country_app, patent_app
        self.date_app, self.country_app, self.patent_app = self.__tagNme__(xmldoc, ["application-reference", ["date", "country", "doc-number"]])
        # code_app
        self.code_app = self.__tagNme__(xmldoc, ["us-application-series-code"])
        # clm_num
        self.clm_num = self.__tagNme__(xmldoc, ["number-of-claims"])
        # classes
        self.classes = [[x[:3].replace(' ',''), x[3:].replace(' ','')] for x in self.__tagNme__(xmldoc, ["classification-national", ["main-classification", "further-classification"]], idx=1, listType=True)]
        # abstract
        self.abstract = self.__allHTML__(xmldoc, ["abstract", "p"])
        # invention_title
        self.invention_title = self.__allHTML__(xmldoc, ["invention-title"])
        # asg_list
        self.asg_list = self.__asg_detail__(self.__tagNme__(xmldoc, ["assignees", "assignee"], iHTML=False))
        # cit_list
        self.cit_list = self.__cit_detail__(self.__tagNme__(xmldoc, ["references-cited", "citation"], iHTML=False))
        # rel_list
        self.rel_list = self.__rel_detail__(self.__tagNme__(xmldoc, ["us-related-documents"], iHTML=False))
        # inv_list
        self.inv_list = self.__tagSplit__(xmldoc, ["parties", "applicant"], [["addressbook", ["last-name", "first-name"]], ["addressbook", "address", ["street", "city", "state", "country", "postcode"]], [["nationality", "residence"], "country"]], blank=True)
        # law_list
        self.law_list = self.__tagSplit__(xmldoc, ["parties", "agents", "agent"], [["addressbook", ["last-name", "first-name", "country", "orgname"]]], blank=True)

    """
    These are all methods from the old XMLPatent class. I'm keeping them
    in here to assist in the carry-over of parsing the xml fields
    """
    def __allHTML__(self, xmldoc, tagList):
        for x in self.__tagNme__(xmldoc, tagList, iHTML=False):
            return re.sub(r"<[/]?%s( .*?)?>" % (tagList[-1]), "", x.toxml())
        return ""

    def __innerHTML__(self, dom_element):
        """
        Gets the text out of all child nodes for a certain
        dom_element
        """
        #if blank return nothing as well!
        if dom_element == '':
            return ''
        else:
            rc = ""
            for node in dom_element.childNodes:
                if node.nodeType == node.TEXT_NODE:
                    rc = rc + node.data
            return rc

    def __tagSplit__(self, xmldoc, xmlList, tagList, baseList=[], idx=0, blank=False, iHTML=True, debug=False):
        d_list = []
        for x in self.__tagNme__(xmldoc, tagList=xmlList, iHTML=False):
            record = copy.copy(baseList)
            for y in tagList:
                if debug:
                    print "-------------------"
                    print x.toxml()
                    print y
                record.extend(self.__tagNme__(x, tagList=y, blank=blank, iHTML=iHTML, debug=debug, idx=idx))
            d_list.append(record)
        return d_list

    def __tagNme__(self, xmldoc, tagList, idx=0, listType=False, blank=False, iHTML=True, debug=False):
        """
        takes in the xml.minidom parsed xml document [xmldoc] and a list of tags [tagList]
        [tagList] may sometimes be a list of lists?
        We iterate the taglist, searching another level deeper into the xml doc with each
        iteration. In tagList, index 0 is the first level of the xmldoc, index 1 is the list
        of all tags we look for at the level right under all the tags we found at index 0,
        and so on.
        For each item, we then iterate through all xmldocs we passed in (usually one?)
        and append the tag we wanted to an xmllist.
        We have empty strings if the tag was not found, and always return a flat list.
        Returns:
        if the iHTML flag is set, we do the innerHTML method on the xmllist, otherwise
        we just return the xmllist.
        """
        xmldoc = [xmldoc]
        for i,x in enumerate(tagList):
            if type(x) is not ListType:
                x = [x]
            xmlnext = []
            for y in x:
                for z in xmldoc:
                    if z != '':
                        if i==0 and idx!=0:
                            if len(z.getElementsByTagName(y))>0:
                                xmlnext.append(z.getElementsByTagName(y)[idx-1])
                            else:
                                xmlnext.append('')
                        else:
                            blFlag = False
                            for za in z.getElementsByTagName(y):
                                blFlag = True
                                xmlnext.append(za)
                            if blFlag==False and blank==True:
                                xmlnext.append('')
                    else:
                        xmlnext.append('')
            xmldoc = xmlnext

        if debug:
            if len(xmldoc)==1 and iHTML:
                print self.__innerHTML__(xmldoc[0])
            elif iHTML:
                print [self.__innerHTML__(x) for x in xmldoc]
            else:
                print xmldoc

        if len(xmldoc)==1 and iHTML:
            if listType:
                return [self.__innerHTML__(xmldoc[0])]
            else:
                return self.__innerHTML__(xmldoc[0])
        elif iHTML:
            return [self.__innerHTML__(x) for x in xmldoc]
        else:
            return xmldoc

    def __asg_detail__(self, xmldoc):
        d_list = []
        for x in xmldoc:
            record = []
            if len(x.getElementsByTagName("first-name"))>0:
                record = [1]
                record.extend(self.__tagNme__(x, [["last-name", "first-name"]]))
            else:
                record = [0]
                record.extend(self.__tagNme__(x, [["orgname", "role"]]))
            record.extend(self.__tagNme__(x, ["addressbook", "address", ["street", "city", "state", "country", "postcode"]], blank=True))
            record.extend(self.__tagNme__(x, [["nationality", "residence"], "country"], blank=True))
            d_list.append(record)
        return d_list

    def __cit_detail__(self, xmldoc):
        d_list = []
        for x in xmldoc:
            #this means patcit is part of the XML
            record = [self.__tagNme__(x, ["category"])]
            if len(x.getElementsByTagName("patcit"))>0:
                record.extend(self.__tagNme__(x, ["patcit", ["country", "doc-number", "date", "kind", "name"]], blank=True))
                record.extend([""])
            elif len(x.getElementsByTagName("othercit"))>0:
                record.extend(["", "", "", "", ""])
                record.extend([self.__allHTML__(x, ["othercit"])])
                #probably should grab date information
            else:
                print x.toxml()
            d_list.append(record)
        return d_list

    def __rel_detail__(self, xmldoc, debug=False):
        d_list = []
        for x in xmldoc:
            for y in ["continuation-in-part", "continuation", "division", "reissue"]:
                if len(x.getElementsByTagName(y))>0:
                    d_list.extend(self.__tagSplit__(x, ["relation", "child-doc"], [[["doc-number", "country", "kind"]]], baseList=[y, -1], blank=True))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc"], [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True, idx=1))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc", "parent-grant-document"], [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc", "parent-pct-document"],  [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True))
            for y in ["related-publication", "us-provisional-application"]:
                if len(x.getElementsByTagName(y))>0:
                    d_list.extend(self.__tagSplit__(x, ["document-id"], [[["doc-number", "country", "kind"]]], baseList=[y, 0], blank=True))
            if debug:
                print "-------------------"
                for x in d_list:
                    print x
        return d_list

    def __repr__(self):
        return \
"""country = %s, patent = %s, pat_type = %s,
date_grant = %s, date_app = %s,
         abstract = %s
  invention_title = %s
len(classes, asg, cit, ret, inv, law) = %s""" % (self.country, self.patent, self.pat_type, self.date_grant, self.date_app,
       self.abstract[:50], self.invention_title[:50],
       str([len(self.classes),  len(self.asg_list), len(self.cit_list),
            len(self.rel_list), len(self.inv_list), len(self.law_list)]))
<commit_msg>add more clarifying comments, still first hash<commit_after>from xml.dom import minidom
from types import *
import datetime
import csv
import os
import re
import sqlite3
import copy
import unicodedata

class XMLPatent(object):
    """
    Base class for handling parsing of xml fields
    """

    def __init__(self, xmlstring):
        """
        [xmlstring]: string containing xml document. Should be provided by
        parse.py
        """
        # lowercase all tags
        self.xmlstring = re.sub(r"<[/]?[A-Za-z-]+?[ >]", lambda x: x.group().lower(), xmlstring)
        # store the minidom parsed xml doc
        self.xmldoc = minidom.parseString(self.xmlstring)

        # country, patent, kind, date_grant
        self.country, self.patent, self.kind, self.date_grant = self.__tagNme__(xmldoc, ["publication-reference", ["country", "doc-number", "kind", "date"]])
        # pat_type
        self.pat_type = self.__tagNme__(xmldoc, ["application-reference"], iHTML=False)[0].attributes["appl-type"].value
        # date_app, country_app, patent_app
        self.date_app, self.country_app, self.patent_app = self.__tagNme__(xmldoc, ["application-reference", ["date", "country", "doc-number"]])
        # code_app
        self.code_app = self.__tagNme__(xmldoc, ["us-application-series-code"])
        # clm_num
        self.clm_num = self.__tagNme__(xmldoc, ["number-of-claims"])
        # classes
        self.classes = [[x[:3].replace(' ',''), x[3:].replace(' ','')] for x in self.__tagNme__(xmldoc, ["classification-national", ["main-classification", "further-classification"]], idx=1, listType=True)]
        # abstract
        self.abstract = self.__allHTML__(xmldoc, ["abstract", "p"])
        # invention_title
        self.invention_title = self.__allHTML__(xmldoc, ["invention-title"])
        # asg_list
        self.asg_list = self.__asg_detail__(self.__tagNme__(xmldoc, ["assignees", "assignee"], iHTML=False))
        # cit_list
        self.cit_list = self.__cit_detail__(self.__tagNme__(xmldoc, ["references-cited", "citation"], iHTML=False))
        # rel_list
        self.rel_list = self.__rel_detail__(self.__tagNme__(xmldoc, ["us-related-documents"], iHTML=False))
        # inv_list
        self.inv_list = self.__tagSplit__(xmldoc, ["parties", "applicant"], [["addressbook", ["last-name", "first-name"]], ["addressbook", "address", ["street", "city", "state", "country", "postcode"]], [["nationality", "residence"], "country"]], blank=True)
        # law_list
        self.law_list = self.__tagSplit__(xmldoc, ["parties", "agents", "agent"], [["addressbook", ["last-name", "first-name", "country", "orgname"]]], blank=True)

    """
    These are all methods from the old XMLPatent class. I'm keeping them
    in here to assist in the carry-over of parsing the xml fields
    """
    def __allHTML__(self, xmldoc, tagList):
        """
        Replaces the html tags for everything returned by __tagName__(taglist)
        with nothing
        """
        for x in self.__tagNme__(xmldoc, tagList, iHTML=False):
            return re.sub(r"<[/]?%s( .*?)?>" % (tagList[-1]), "", x.toxml())
        return ""

    def __innerHTML__(self, dom_element):
        """
        Gets the text out of all child nodes for a certain
        dom_element
        """
        #if blank return nothing as well!
        if dom_element == '':
            return ''
        else:
            rc = ""
            for node in dom_element.childNodes:
                if node.nodeType == node.TEXT_NODE:
                    rc = rc + node.data
            return rc

    def __tagSplit__(self, xmldoc, xmlList, tagList, baseList=[], idx=0, blank=False, iHTML=True, debug=False):
        """
        For each item in the list of xml tags returned by tagName, we treat each tag's contents
        as its own xml doc and search it for a list of xml tags. Everyything is appended together
        and returned
        """
        d_list = []
        for x in self.__tagNme__(xmldoc, tagList=xmlList, iHTML=False):
            record = copy.copy(baseList)
            for y in tagList:
                if debug:
                    print "-------------------"
                    print x.toxml()
                    print y
                record.extend(self.__tagNme__(x, tagList=y, blank=blank, iHTML=iHTML, debug=debug, idx=idx))
            d_list.append(record)
        return d_list

    def __tagNme__(self, xmldoc, tagList, idx=0, listType=False, blank=False, iHTML=True, debug=False):
        """
        takes in the xml.minidom parsed xml document [xmldoc] and a list of tags [tagList]
        [tagList] may sometimes be a list of lists?
        We iterate the taglist, searching another level deeper into the xml doc with each
        iteration. In tagList, index 0 is the first level of the xmldoc, index 1 is the list
        of all tags we look for at the level right under all the tags we found at index 0,
        and so on.
        For each item, we then iterate through all xmldocs we passed in (usually one?)
        and append the tag we wanted to an xmllist.
        We have empty strings if the tag was not found, and always return a flat list.
        **Returns:
        if the iHTML flag is set, we do the innerHTML method on the xmllist, otherwise
        we just return the xmllist.
        """
        xmldoc = [xmldoc]
        for i,x in enumerate(tagList):
            if type(x) is not ListType:
                x = [x]
            xmlnext = []
            for y in x:
                for z in xmldoc:
                    if z != '':
                        if i==0 and idx!=0:
                            if len(z.getElementsByTagName(y))>0:
                                xmlnext.append(z.getElementsByTagName(y)[idx-1])
                            else:
                                xmlnext.append('')
                        else:
                            blFlag = False
                            for za in z.getElementsByTagName(y):
                                blFlag = True
                                xmlnext.append(za)
                            if blFlag==False and blank==True:
                                xmlnext.append('')
                    else:
                        xmlnext.append('')
            xmldoc = xmlnext

        if debug:
            if len(xmldoc)==1 and iHTML:
                print self.__innerHTML__(xmldoc[0])
            elif iHTML:
                print [self.__innerHTML__(x) for x in xmldoc]
            else:
                print xmldoc

        if len(xmldoc)==1 and iHTML:
            if listType:
                return [self.__innerHTML__(xmldoc[0])]
            else:
                return self.__innerHTML__(xmldoc[0])
        elif iHTML:
            return [self.__innerHTML__(x) for x in xmldoc]
        else:
            return xmldoc

    def __asg_detail__(self, xmldoc):
        d_list = []
        for x in xmldoc:
            record = []
            if len(x.getElementsByTagName("first-name"))>0:
                record = [1]
                record.extend(self.__tagNme__(x, [["last-name", "first-name"]]))
            else:
                record = [0]
                record.extend(self.__tagNme__(x, [["orgname", "role"]]))
            record.extend(self.__tagNme__(x, ["addressbook", "address", ["street", "city", "state", "country", "postcode"]], blank=True))
            record.extend(self.__tagNme__(x, [["nationality", "residence"], "country"], blank=True))
            d_list.append(record)
        return d_list

    def __cit_detail__(self, xmldoc):
        d_list = []
        for x in xmldoc:
            #this means patcit is part of the XML
            record = [self.__tagNme__(x, ["category"])]
            if len(x.getElementsByTagName("patcit"))>0:
                record.extend(self.__tagNme__(x, ["patcit", ["country", "doc-number", "date", "kind", "name"]], blank=True))
                record.extend([""])
            elif len(x.getElementsByTagName("othercit"))>0:
                record.extend(["", "", "", "", ""])
                record.extend([self.__allHTML__(x, ["othercit"])])
                #probably should grab date information
            else:
                print x.toxml()
            d_list.append(record)
        return d_list

    def __rel_detail__(self, xmldoc, debug=False):
        d_list = []
        for x in xmldoc:
            for y in ["continuation-in-part", "continuation", "division", "reissue"]:
                if len(x.getElementsByTagName(y))>0:
                    d_list.extend(self.__tagSplit__(x, ["relation", "child-doc"], [[["doc-number", "country", "kind"]]], baseList=[y, -1], blank=True))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc"], [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True, idx=1))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc", "parent-grant-document"], [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True))
                    d_list.extend(self.__tagSplit__(x, ["relation", "parent-doc", "parent-pct-document"],  [[["doc-number", "country", "kind", "date", "parent-status"]]], baseList=[y, 1], blank=True))
            for y in ["related-publication", "us-provisional-application"]:
                if len(x.getElementsByTagName(y))>0:
                    d_list.extend(self.__tagSplit__(x, ["document-id"], [[["doc-number", "country", "kind"]]], baseList=[y, 0], blank=True))
            if debug:
                print "-------------------"
                for x in d_list:
                    print x
        return d_list

    def __repr__(self):
        return \
"""country = %s, patent = %s, pat_type = %s,
date_grant = %s, date_app = %s,
         abstract = %s
  invention_title = %s
len(classes, asg, cit, ret, inv, law) = %s""" % (self.country, self.patent, self.pat_type, self.date_grant, self.date_app,
       self.abstract[:50], self.invention_title[:50],
       str([len(self.classes),  len(self.asg_list), len(self.cit_list),
            len(self.rel_list), len(self.inv_list), len(self.law_list)]))
<|endoftext|>


<commit_before>import sys, os, shutil

import inv_check
import insert_jaif
import ontology_to_daikon

import backend
import common

def run_inference(project):
  project_dir = common.get_project_dir(project)
  annotation_dir = os.path.join(project_dir, common.DLJC_OUTPUT_DIR, 'annotations')

  jsr308 = common.TOOLS_DIR
  os.environ['JSR308'] = jsr308
  classpath = os.path.join(jsr308, 'generic-type-inference-solver', 'bin')
  if os.environ.get('CLASSPATH'):
    os.environ['CLASSPATH'] += ':' + classpath
  else:
    os.environ['CLASSPATH'] = classpath

  afu = os.path.join(jsr308, 'annotation-tools', 'annotation-file-utilities')
  os.environ['AFU'] = afu
  os.environ['PATH'] += ':' + os.path.join(afu, 'scripts')

  if os.path.isdir(annotation_dir):
    shutil.rmtree(annotation_dir)

  with common.cd(project_dir):
    common.clean_project(project)
    common.run_dljc(project,
                    ['inference'],
                    ['--solverArgs=backEndType=maxsatbackend.MaxSat',
                     '--checker', 'ontology.OntologyChecker',
                     '--solver', 'constraintsolver.ConstraintSolver',
                     '-m', 'ROUNDTRIP',
                     '-afud', annotation_dir])

def read_jaif_file(jaif_file):
  current_package = ""
  current_class = ""
  current_method = ""
  with open(jaif_file, 'r') as f:
    for line in f.readlines():
      if line.startswith("package "):
        current_package = line[len("package "):line.find(":")]
      if line.startswith("method "):
        current_class = line[len("class "):line.find(":")]
      if line.startswith("method "):
        current_method = line[len("method "):line.find(":")]

      if line.startswith("insert-annotation Method.parameter"):
        print ("YEAH {}{}".format(current_class, current_method))
      if line.startswith("insert-annotation Method.type"):
        print ("RET {}{}".format(current_class, current_method))

def find_methods_with_signature(corpus, return_annotation, param_annotation_list):

  good_methods = []

  for project in corpus:
    project_dir = common.get_project_dir(project)
    jaif_file = os.path.join(project_dir, "default.jaif")

    has_param = False
    has_ret = False
    current_package = ""
    current_class = ""
    current_method = ""
    with open(jaif_file, 'r') as f:
      for line in f.readlines():
        if line.startswith("package "):
          current_package = line[len("package "):line.find(":")]
        if line.startswith("class "):
          current_class = line[len("class "):line.find(":")]
          has_param = False
          has_ret = False
        if line.startswith("method "):
          current_method = line[len("method "):line.find(":")]

        if line.startswith("insert-annotation Method.parameter"):
          s = line[len("insert-annotation Method.parameter "):]
          param_idx = int(s[:s.find(",")])
          if len(param_annotation_list) > param_idx and param_annotation_list[param_idx] in line:
            has_param = True
          else:
            has_param = False
        if line.startswith("insert-annotation Method.type") and return_annotation in line:
          has_ret = True

        if has_param==True and has_ret==True:
          good_methods += [(project, current_package, current_class, current_method)]
          print ("Relevant Method: {}.{}".format(current_class,current_method))
          has_param = False
          has_ret = False

def main(corpus):
  """ SUMMARY: use case of the user-driven functionality of PASCALI.
  Scenario: User provides the concept of Sequence and the equivalent Java
  types, and the concept of sorted sequence and the relevant type invariant.
  Goal: learn how to get from Sequence -> Sorted Sequence.
  """

  """ Look for new mapping from 'ontology concepts'->'java type' and run
  checker framework. Should be implemented in type_inference
  Mapping example:
    Sequence -> java.lang.Array, java.util.List, LinkedHashSet, etc.

  INPUT: corpus, file containing set of concept->java_type mapping
  OUTPUT: Set of jaif files that are merged into the classes using
          insert_jaif.
  BODY: This also triggers back-end labeled graph generation.
  """

  for project in corpus:
    run_inference(project)

  """ Missing step: interact with PA to add a definition of Sorted Sequence
  which is a specialization of Sequence that has a sortedness invariants.
  The sortedness invariant gets turned into a Daikon template
  INPUT: user interaction
  OUTPUT: type_annotation and type_invariant (for sorted sequence)

  """

  ontology_invariant_file = "TODO_from_Howie.txt"
  with open(ontology_invariant_file, 'w') as f:
    f.write("TODO")

  invariant_name = "TODO_sorted_sequence"

  daikon_pattern_java_file = ontology_to_daikon.create_daikon_invariant(ontology_invariant_file, invariant_name)

  """ Search for methods that have a return type annotated with Sequence
  and for which we can establish a sortedness invariant (may done by LB).

  INPUT: dtrace file of project
         daikon_pattern_java_file that we want to check on the dtrace file.

  OUTPUT: list of ppt names that establish the invariant. Here a ppt
  is a Daikon program point, s.a. test01.TestClass01.sort(int[]):::EXIT

  Note: this step translate the type_invariant into a Daikon
  template (which is a Java file).
  """

  pattern_class_name = invariant_name
  pattern_class_dir = os.path.join(common.WORKING_DIR, "invClass")
  if os.path.isdir(pattern_class_dir):
    shutil.rmtree(pattern_class_dir)
  os.mkdir(pattern_class_dir)

  cmd = ["javac", "-g", "-classpath", common.get_jar('daikon.jar'),
         daikon_pattern_java_file, "-d", pattern_class_dir]
  common.run_cmd(cmd)

  list_of_methods = []
  for project in corpus:
    dtrace_file = backend.get_dtrace_file_for_project(project)
    if not dtrace_file:
      print ("Ignoring folder {} because it does not contain dtrace file".format(project))
      continue
    ppt_names = inv_check.find_ppts_that_establish_inv(dtrace_file, pattern_class_dir, pattern_class_name)
    methods = set()
    for ppt in ppt_names:
      method_name = ppt[:ppt.find(':::EXIT')]
      methods.add(method_name)
    list_of_methods +=[(project, methods)]

  print ("\n   ************")
  print ("The following corpus methods establish Howies invariant in the return value:")
  for project, methods in list_of_methods:
    if len(methods)>0:
      print (project)
      for m in methods:
        print("\t{}".format(m))
  print ("\n   ************")

  """ Expansion of dynamic analysis results ....
  Find a list of similar methods that are similar to the ones found above (list_of_methods).
  INPUT: list_of_methods, corpus with labeled graphs generated, threshold value for similarity,
  OUTPUT: superset_list_of_methods
  """

  # WENCHAO
  print ("Expanding the dynamic analysis results using graph-based similarity:")
  union_set = set()
  for project, methods in list_of_methods:
    # map Daikon output on sort method to method signature in methods.txt in generated graphs
    for m in methods:
      method_name = common.get_method_from_daikon_out(m)
      #kernel_file = common.get_kernel_path(project)
      method_file = common.get_method_path(project)
      dot_name = common.find_dot_name(method_name, method_file)
      if dot_name:
        # find the right dot file for each method
        dot_file = common.get_dot_path(project, dot_name)
        # find all graphs that are similar to it using WL based on some threshold
        sys.path.insert(0, 'simprog')
        from similarity import Similarity
        sim = Similarity()
        sim.read_graph_kernels("corpus_kernel.txt")
        top_k = 3
        iter_num = 3
        result_program_list_with_score = sim.find_top_k_similar_graphs(dot_file, 'g', top_k, iter_num)
        result_set = set([x[0] for x in result_program_list_with_score])
        # take the union of all these graphs
        union_set = union_set | result_set
  print union_set

  # return this set as a list of (project, method)
  fo = open("methods.txt", "w")
  expanded_list = []
  for dot_path in union_set:
    method_summary = common.get_method_summary_from_dot_path(dot_path)
    fo.write(method_summary)
    fo.write("\n")
  fo.close()

  """ Update the type annotations for the expanded dynamic analysis results.
  INPUT: superset_list_of_methods, annotation to be added
  OUTPUT: nothing
  EFFECT: updates the type annotations of the methods in superset_list_of_methods.
  This requires some additional checks to make sure that the methods actually
  perform some kind of sorting. Note that we do it on the superset because the original
  list_of_methods might miss many implementations because fuzz testing could not
  reach them.
  """
  for class_file in []: # MARTIN
    generated_jaif_file = "TODO"
    insert_jaif.merge_jaif_into_class(class_file, generated_jaif_file)

  """ Ordering of expanded dynamic analysis results ....
  Find the k 'best' implementations in superset of list_of_methods
  INPUT: superset_list_of_methods, corpus, k
  OUTPUT: k_list_of_methods
  Note: similarity score is used. may consider using other scores; e.g., TODO:???
  """

  #TODO: create input file for huascar where each line is formatted like:
  # ../corpus/Sort05/src/Sort05.java::sort(int[]):int[]

  ordering_dir = os.path.join(common.WORKING_DIR, "ordering_results/")

  with common.cd(ordering_dir):
    #TODO generate a proper relevant methods file.
    relevant_methods_file = "methods.txt"
    cmd = ["./cue",
           "typical",
           "-k", "3",
           "-f", relevant_methods_file]
    common.run_cmd(cmd, print_output=True)

  """
  Close the loop and add the best implementation found in the previous
  step back to the ontology.
  INPUT: k_list_of_methods
  OUTPUT: patch file for the ontology. Worst case: just add the 'best' implementation
  found in the corpus as a blob to the ontology. Best case: generate an equivalent
  flow-graph in the ontology.
  """
  print "TODO" # ALL

# if not os.path.isfile(daikon_jar):
#   print "Downloading dependencies"
#   cmd = ["./fetch_dependencies.sh"]
#   run_command(cmd)
#   print "Done."

if __name__ == '__main__':
  corpus = common.get_project_list()
  if len(sys.argv)>1:
    filtered_corpus = []
    for arg in sys.argv[1:]:
      if arg in corpus:
        filtered_corpus += [arg]
    print ("Filtered corpus contianing: {}".format(','.join(filtered_corpus)))
    corpus = filtered_corpus
  main(corpus)
  #find_methods_with_signature(corpus, "@ontology.qual.Sequence", ["@ontology.qual.Sequence"])

<commit_msg>better print message for expanded set of programs<commit_after>import sys, os, shutil

import inv_check
import insert_jaif
import ontology_to_daikon

import backend
import common

def run_inference(project):
  project_dir = common.get_project_dir(project)
  annotation_dir = os.path.join(project_dir, common.DLJC_OUTPUT_DIR, 'annotations')

  jsr308 = common.TOOLS_DIR
  os.environ['JSR308'] = jsr308
  classpath = os.path.join(jsr308, 'generic-type-inference-solver', 'bin')
  if os.environ.get('CLASSPATH'):
    os.environ['CLASSPATH'] += ':' + classpath
  else:
    os.environ['CLASSPATH'] = classpath

  afu = os.path.join(jsr308, 'annotation-tools', 'annotation-file-utilities')
  os.environ['AFU'] = afu
  os.environ['PATH'] += ':' + os.path.join(afu, 'scripts')

  if os.path.isdir(annotation_dir):
    shutil.rmtree(annotation_dir)

  with common.cd(project_dir):
    common.clean_project(project)
    common.run_dljc(project,
                    ['inference'],
                    ['--solverArgs=backEndType=maxsatbackend.MaxSat',
                     '--checker', 'ontology.OntologyChecker',
                     '--solver', 'constraintsolver.ConstraintSolver',
                     '-m', 'ROUNDTRIP',
                     '-afud', annotation_dir])

def read_jaif_file(jaif_file):
  current_package = ""
  current_class = ""
  current_method = ""
  with open(jaif_file, 'r') as f:
    for line in f.readlines():
      if line.startswith("package "):
        current_package = line[len("package "):line.find(":")]
      if line.startswith("method "):
        current_class = line[len("class "):line.find(":")]
      if line.startswith("method "):
        current_method = line[len("method "):line.find(":")]

      if line.startswith("insert-annotation Method.parameter"):
        print ("YEAH {}{}".format(current_class, current_method))
      if line.startswith("insert-annotation Method.type"):
        print ("RET {}{}".format(current_class, current_method))

def find_methods_with_signature(corpus, return_annotation, param_annotation_list):

  good_methods = []

  for project in corpus:
    project_dir = common.get_project_dir(project)
    jaif_file = os.path.join(project_dir, "default.jaif")

    has_param = False
    has_ret = False
    current_package = ""
    current_class = ""
    current_method = ""
    with open(jaif_file, 'r') as f:
      for line in f.readlines():
        if line.startswith("package "):
          current_package = line[len("package "):line.find(":")]
        if line.startswith("class "):
          current_class = line[len("class "):line.find(":")]
          has_param = False
          has_ret = False
        if line.startswith("method "):
          current_method = line[len("method "):line.find(":")]

        if line.startswith("insert-annotation Method.parameter"):
          s = line[len("insert-annotation Method.parameter "):]
          param_idx = int(s[:s.find(",")])
          if len(param_annotation_list) > param_idx and param_annotation_list[param_idx] in line:
            has_param = True
          else:
            has_param = False
        if line.startswith("insert-annotation Method.type") and return_annotation in line:
          has_ret = True

        if has_param==True and has_ret==True:
          good_methods += [(project, current_package, current_class, current_method)]
          print ("Relevant Method: {}.{}".format(current_class,current_method))
          has_param = False
          has_ret = False

def main(corpus):
  """ SUMMARY: use case of the user-driven functionality of PASCALI.
  Scenario: User provides the concept of Sequence and the equivalent Java
  types, and the concept of sorted sequence and the relevant type invariant.
  Goal: learn how to get from Sequence -> Sorted Sequence.
  """

  """ Look for new mapping from 'ontology concepts'->'java type' and run
  checker framework. Should be implemented in type_inference
  Mapping example:
    Sequence -> java.lang.Array, java.util.List, LinkedHashSet, etc.

  INPUT: corpus, file containing set of concept->java_type mapping
  OUTPUT: Set of jaif files that are merged into the classes using
          insert_jaif.
  BODY: This also triggers back-end labeled graph generation.
  """

  for project in corpus:
    run_inference(project)

  """ Missing step: interact with PA to add a definition of Sorted Sequence
  which is a specialization of Sequence that has a sortedness invariants.
  The sortedness invariant gets turned into a Daikon template
  INPUT: user interaction
  OUTPUT: type_annotation and type_invariant (for sorted sequence)

  """

  ontology_invariant_file = "TODO_from_Howie.txt"
  with open(ontology_invariant_file, 'w') as f:
    f.write("TODO")

  invariant_name = "TODO_sorted_sequence"

  daikon_pattern_java_file = ontology_to_daikon.create_daikon_invariant(ontology_invariant_file, invariant_name)

  """ Search for methods that have a return type annotated with Sequence
  and for which we can establish a sortedness invariant (may done by LB).

  INPUT: dtrace file of project
         daikon_pattern_java_file that we want to check on the dtrace file.

  OUTPUT: list of ppt names that establish the invariant. Here a ppt
  is a Daikon program point, s.a. test01.TestClass01.sort(int[]):::EXIT

  Note: this step translate the type_invariant into a Daikon
  template (which is a Java file).
  """

  pattern_class_name = invariant_name
  pattern_class_dir = os.path.join(common.WORKING_DIR, "invClass")
  if os.path.isdir(pattern_class_dir):
    shutil.rmtree(pattern_class_dir)
  os.mkdir(pattern_class_dir)

  cmd = ["javac", "-g", "-classpath", common.get_jar('daikon.jar'),
         daikon_pattern_java_file, "-d", pattern_class_dir]
  common.run_cmd(cmd)

  list_of_methods = []
  for project in corpus:
    dtrace_file = backend.get_dtrace_file_for_project(project)
    if not dtrace_file:
      print ("Ignoring folder {} because it does not contain dtrace file".format(project))
      continue
    ppt_names = inv_check.find_ppts_that_establish_inv(dtrace_file, pattern_class_dir, pattern_class_name)
    methods = set()
    for ppt in ppt_names:
      method_name = ppt[:ppt.find(':::EXIT')]
      methods.add(method_name)
    list_of_methods +=[(project, methods)]

  print ("\n   ************")
  print ("The following corpus methods establish Howies invariant in the return value:")
  for project, methods in list_of_methods:
    if len(methods)>0:
      print (project)
      for m in methods:
        print("\t{}".format(m))
  print ("\n   ************")

  """ Expansion of dynamic analysis results ....
  Find a list of similar methods that are similar to the ones found above (list_of_methods).
  INPUT: list_of_methods, corpus with labeled graphs generated, threshold value for similarity,
  OUTPUT: superset_list_of_methods
  """

  # WENCHAO
  print("Expanding the dynamic analysis results using graph-based similarity:")
  union_set = set()
  for project, methods in list_of_methods:
    # map Daikon output on sort method to method signature in methods.txt in generated graphs
    for m in methods:
      method_name = common.get_method_from_daikon_out(m)
      #kernel_file = common.get_kernel_path(project)
      method_file = common.get_method_path(project)
      dot_name = common.find_dot_name(method_name, method_file)
      if dot_name:
        # find the right dot file for each method
        dot_file = common.get_dot_path(project, dot_name)
        # find all graphs that are similar to it using WL based on some threshold
        sys.path.insert(0, 'simprog')
        from similarity import Similarity
        sim = Similarity()
        sim.read_graph_kernels("corpus_kernel.txt")
        top_k = 3
        iter_num = 3
        result_program_list_with_score = sim.find_top_k_similar_graphs(dot_file, 'g', top_k, iter_num)
        result_set = set([x[0] for x in result_program_list_with_score])
        # take the union of all these graphs
        union_set = union_set | result_set
  print("Expanded set:")
  print([x.split('/')[-4] for x in union_set])

  # return this set as a list of (project, method)
  fo = open("methods.txt", "w")
  expanded_list = []
  for dot_path in union_set:
    method_summary = common.get_method_summary_from_dot_path(dot_path)
    fo.write(method_summary)
    fo.write("\n")
  fo.close()

  """ Update the type annotations for the expanded dynamic analysis results.
  INPUT: superset_list_of_methods, annotation to be added
  OUTPUT: nothing
  EFFECT: updates the type annotations of the methods in superset_list_of_methods.
  This requires some additional checks to make sure that the methods actually
  perform some kind of sorting. Note that we do it on the superset because the original
  list_of_methods might miss many implementations because fuzz testing could not
  reach them.
  """
  for class_file in []: # MARTIN
    generated_jaif_file = "TODO"
    insert_jaif.merge_jaif_into_class(class_file, generated_jaif_file)

  """ Ordering of expanded dynamic analysis results ....
  Find the k 'best' implementations in superset of list_of_methods
  INPUT: superset_list_of_methods, corpus, k
  OUTPUT: k_list_of_methods
  Note: similarity score is used. may consider using other scores; e.g., TODO:???
  """

  #TODO: create input file for huascar where each line is formatted like:
  # ../corpus/Sort05/src/Sort05.java::sort(int[]):int[]

  ordering_dir = os.path.join(common.WORKING_DIR, "ordering_results/")

  with common.cd(ordering_dir):
    #TODO generate a proper relevant methods file.
    relevant_methods_file = "methods.txt"
    cmd = ["./cue",
           "typical",
           "-k", "3",
           "-f", relevant_methods_file]
    common.run_cmd(cmd, print_output=True)

  """
  Close the loop and add the best implementation found in the previous
  step back to the ontology.
  INPUT: k_list_of_methods
  OUTPUT: patch file for the ontology. Worst case: just add the 'best' implementation
  found in the corpus as a blob to the ontology. Best case: generate an equivalent
  flow-graph in the ontology.
  """
  print "TODO" # ALL

# if not os.path.isfile(daikon_jar):
#   print "Downloading dependencies"
#   cmd = ["./fetch_dependencies.sh"]
#   run_command(cmd)
#   print "Done."

if __name__ == '__main__':
  corpus = common.get_project_list()
  if len(sys.argv)>1:
    filtered_corpus = []
    for arg in sys.argv[1:]:
      if arg in corpus:
        filtered_corpus += [arg]
    print ("Filtered corpus contianing: {}".format(','.join(filtered_corpus)))
    corpus = filtered_corpus
  main(corpus)
  #find_methods_with_signature(corpus, "@ontology.qual.Sequence", ["@ontology.qual.Sequence"])

<|endoftext|>


<commit_before>import os
import struct
import logging

from primitives import *
from constants import *
from astypes import MalformedFLV
from astypes import get_script_data_variable, make_script_data_variable

log = logging.getLogger('flvlib.tags')

STRICT_PARSING = False
def strict_parser():
    return globals()['STRICT_PARSING']

class EndOfTags(Exception):
    pass

def ensure(value, expected, error_msg):
    if value == expected:
        return

    if strict_parser():
        raise MalformedFLV(error_msg)
    else:
        log.warning('Skipping non-conformant value in FLV file')

class Tag(object):

    def __init__(self, parent_flv, f):
        self.f = f
        self.parent_flv = parent_flv
        self.offset = None
        self.size = None
        self.timestamp = None
        self.data_start = None

    def parse(self):
        f = self.f

        self.offset = f.tell() - 1

        # DataSize
        self.size = get_ui24(f)
        log.debug("Data size is %d", self.size)

        # Timestamp + TimestampExtended
        self.timestamp = get_si32_extended(f)
        log.debug("Tag timestamp is %d", self.timestamp)

        if self.timestamp < 0:
            log.warning("The tag at offset 0x%08X has negative timestamp: %d",
                        self.offset, self.timestamp)

        # StreamID
        stream_id = get_ui24(f)
        ensure(stream_id, 0, "StreamID non zero: 0x%06X" % stream_id)

        # The rest gets parsed in the subclass

    def seek_past_data(self, adjust=0):
        f = self.f

        f.seek(self.size + adjust, os.SEEK_CUR)
        log.debug("%s seeked past the data of the tag", self)

        previous_tag_size = get_ui32(f)
        ensure(previous_tag_size, self.size + 11,
               "PreviousTagSize of %d (0x%08X) "
               "not equal to actual tag size of %d (0x%08X)" %
               (previous_tag_size, previous_tag_size,
                self.size + 11, self.size + 11))
        log.debug("Ready to read another tag")

class AudioTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.sound_format = None
        self.sound_rate = None
        self.sound_size = None
        self.sound_type = None

    def parse(self):
        f = self.f
        Tag.parse(self)

        sound_flags = get_ui8(f)

        self.sound_format = (sound_flags & 0xF0) >> 4
        self.sound_rate = (sound_flags & 0xC) >> 2
        self.sound_size = (sound_flags & 0x2) >> 1
        self.sound_type = sound_flags & 0x1

        if strict_parser():
            try:
                sound_format_to_string[self.sound_format]
            except KeyError:
                raise MalformedFLV("Invalid sound format: %d",
                                   self.sound_format)
            try:
                sound_rate_to_string[self.sound_rate]
            except KeyError:
                raise MalformedFLV("Invalid sound rate: %d",
                                   self.sound_rate)
            try:
                sound_size_to_string[self.sound_size]
            except KeyError:
                raise MalformedFLV("Invalid sound size type: %d",
                                   self.sound_size)
            try:
                sound_type_to_string[self.sound_type]
            except KeyError:
                raise MalformedFLV("Invalid sound type: %d",
                                   self.sound_type)

        if self.sound_format == SOUND_FORMAT_AAC:
            # AAC always has sampling rate of 44 kHz
            ensure(self.sound_rate, SOUND_RATE_44_KHZ,
                   "AAC sound format with incorrect sound rate: %d" %
                   self.sound_rate)
            # AAC is always stereo
            ensure(self.sound_type, SOUND_TYPE_STEREO,
                   "AAC sound format with incorrect sound type: %d" %
                   self.sound_type)

        self.seek_past_data(adjust=-1)

    def __repr__(self):
        if self.offset is None:
            return "<AudioTag unparsed>"
        elif self.sound_format is None:
            return ("<AudioTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<AudioTag at offset 0x%08X, time %d, size %d, %s>" %
                    (self.offset, self.timestamp, self.size,
                     sound_format_to_string[self.sound_format]))

class VideoTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.frame_type = None
        self.codec_id = None

    def parse(self):
        f = self.f
        Tag.parse(self)

        video_flags = get_ui8(f)

        self.frame_type = (video_flags & 0xF0) >> 4
        self.codec_id = video_flags & 0xF

        if strict_parser():
            try:
                frame_type_to_string[self.frame_type]
            except KeyError:
                raise MalformedFLV("Invalid frame type: %d", self.frame_type)
            try:
                codec_id_to_string[self.codec_id]
            except KeyError:
                raise MalformedFLV("Invalid codec ID: %d", self.codec_id)

        self.seek_past_data(adjust=-1)

    def __repr__(self):
        if self.offset is None:
            return "<VideoTag unparsed>"
        elif self.frame_type is None:
            return ("<VideoTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<VideoTag at offset 0x%08X, time %d, size %d, %s (%s)>" %
                    (self.offset, self.timestamp, self.size,
                     codec_id_to_string[self.codec_id],
                     frame_type_to_string[self.frame_type]))

class ScriptTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.name = None
        self.variable = None

    def parse(self):
        f = self.f
        Tag.parse(self)

        # Here there's always a byte with the value of 0x02,
        # which means "string", although the spec says NOTHING
        # about it..
        value_type = get_ui8(f)
        ensure(value_type, 2, "The name of a script tag is not a string")

        # Need to pass the tag end offset, because apparently YouTube
        # doesn't give a *shit* about the FLV spec and just happily
        # ends the onMetaData tag after self.size bytes, instead of
        # ending it with the *required* 0x09 marker. Bastards!

        if strict_parser():
            # If we're strict, just don't pass this info
            tag_end = None
        else:
            # 11 = tag type (1) + data size (3) + timestamp (4) + stream id (3)
            tag_end = self.offset + 11 + self.size

        log.debug("max offset is 0x%08X", tag_end)

        self.name, self.variable = \
                   get_script_data_variable(f, max_offset=tag_end)
        log.debug("A script tag with a name of %s and value of %r",
                  self.name, self.variable)

        previous_tag_size = get_ui32(f)
        ensure(previous_tag_size, self.size + 11,
               "PreviousTagSize of %d (0x%08X) "
               "not equal to actual tag size of %d (0x%08X)" %
               (previous_tag_size, previous_tag_size,
                self.size + 11, self.size + 11))
        log.debug("Ready to read another tag")

    def __repr__(self):
        if self.offset is None:
            return "<ScriptTag unparsed>"
        elif not self.name:
            return ("<ScriptTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<ScriptTag %s at offset 0x%08X, time %d, size %d>" %
                    (self.name, self.offset, self.timestamp, self.size))

tag_to_class = {
    TAG_TYPE_AUDIO: AudioTag,
    TAG_TYPE_VIDEO: VideoTag,
    TAG_TYPE_SCRIPT: ScriptTag
}

class FLV(object):

    def __init__(self, f):
        self.f = f
        self.version = None
        self.has_audio = None
        self.has_video = None
        self.tags = []

    def parse_header(self):
        f = self.f
        f.seek(0)

        # FLV header
        header = f.read(3)
        if len(header) < 3:
            raise MalformedFLV("The file is shorter than 3 bytes")

        # Do this irrelevant of STRICT_PARSING, to catch bogus files
        if header != "FLV":
            raise MalformedFLV("File signature is incorrect: 0x%X 0x%X 0x%X" %
                               struct.unpack("3B", header))

        # File version
        self.version = get_ui8(f)
        log.debug("File version is %d", self.version)

        # TypeFlags
        flags = get_ui8(f)

        ensure(flags & 0xF8, 0,
               "First TypeFlagsReserved field non zero: 0x%X" % (flags & 0xF8))
        ensure(flags & 0x2, 0,
               "Second TypeFlagsReserved field non zero: 0x%X" % (flags & 0x2))

        self.has_audio = False
        self.has_video = False
        if flags & 0x4:
            self.has_audio = True
        if flags & 0x1:
            self.has_video = True
        log.debug("File %s audio",
                  (self.has_audio and "has") or "does not have")
        log.debug("File %s video",
                  (self.has_video and "has") or "does not have")

        header_size = get_ui32(f)
        log.debug("Header size is %d bytes", header_size)

        f.seek(header_size)

        tag_0_size = get_ui32(f)
        ensure(tag_0_size, 0, "PreviousTagSize0 non zero: 0x%08X" % tag_0_size)

    def iter_tags(self):
        self.parse_header()
        try:
            while True:
                tag = self.get_next_tag()
                yield tag
        except EndOfTags:
            pass

    def read_tags(self):
        self.tags = list(self.iter_tags())

    def get_next_tag(self):
        f = self.f

        try:
            tag_type = get_ui8(f)
        except EndOfFile:
            raise EndOfTags

        tag_klass = self.tag_type_to_class(tag_type)
        tag = tag_klass(self, f)
        log.debug("Found a tag: %s", tag)

        tag.parse()

        return tag

    def tag_type_to_class(self, tag_type):
        try:
            return tag_to_class[tag_type]
        except KeyError:
            raise MalformedFLV("Invalid tag type: %d", tag_type)

def create_flv_tag(type, data, timestamp=0):
    tag_type = struct.pack("B", type)
    timestamp = make_si32_extended(timestamp)
    stream_id = make_ui24(0)

    data_size = len(data)
    tag_size = data_size + 11

    return ''.join([tag_type, make_ui24(data_size), timestamp, stream_id,
                    data, make_ui32(tag_size)])

def create_script_tag(name, data, timestamp=0):
    payload = make_ui8(2) + make_script_data_variable(name, data)
    return create_flv_tag(TAG_TYPE_SCRIPT, payload, timestamp)

def create_flv_header(has_audio=True, has_video=True):
    type_flags = 0
    if has_video:
        type_flags = type_flags | 0x1
    if has_audio:
        type_flags = type_flags | 0x4
    return ''.join(['FLV', make_ui8(1), make_ui8(type_flags), make_ui32(9),
                    make_ui32(0)])
<commit_msg>Split tag parsing into parse() and parse_tag_content() methods.<commit_after>import os
import struct
import logging

from primitives import *
from constants import *
from astypes import MalformedFLV
from astypes import get_script_data_variable, make_script_data_variable

log = logging.getLogger('flvlib.tags')

STRICT_PARSING = False
def strict_parser():
    return globals()['STRICT_PARSING']

class EndOfTags(Exception):
    pass

def ensure(value, expected, error_msg):
    if value == expected:
        return

    if strict_parser():
        raise MalformedFLV(error_msg)
    else:
        log.warning('Skipping non-conformant value in FLV file')

class Tag(object):

    def __init__(self, parent_flv, f):
        self.f = f
        self.parent_flv = parent_flv
        self.offset = None
        self.size = None
        self.timestamp = None
        self.data_start = None

    def parse(self):
        f = self.f

        self.offset = f.tell() - 1

        # DataSize
        self.size = get_ui24(f)
        log.debug("Data size is %d", self.size)

        # Timestamp + TimestampExtended
        self.timestamp = get_si32_extended(f)
        log.debug("Tag timestamp is %d", self.timestamp)

        if self.timestamp < 0:
            log.warning("The tag at offset 0x%08X has negative timestamp: %d",
                        self.offset, self.timestamp)

        # StreamID
        stream_id = get_ui24(f)
        ensure(stream_id, 0, "StreamID non zero: 0x%06X" % stream_id)

        # The rest gets parsed in the subclass, it should move f to the
        # correct position to read PreviousTagSize
        self.parse_tag_content()

        previous_tag_size = get_ui32(f)
        ensure(previous_tag_size, self.size + 11,
               "PreviousTagSize of %d (0x%08X) "
               "not equal to actual tag size of %d (0x%08X)" %
               (previous_tag_size, previous_tag_size,
                self.size + 11, self.size + 11))
        log.debug("Ready to read another tag")

    def parse_tag_content(self):
        # By default just seek past the tag content
        self.f.seek(self.size, os.SEEK_CUR)

class AudioTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.sound_format = None
        self.sound_rate = None
        self.sound_size = None
        self.sound_type = None

    def parse_tag_content(self):
        f = self.f

        sound_flags = get_ui8(f)
        read_bytes = 1

        self.sound_format = (sound_flags & 0xF0) >> 4
        self.sound_rate = (sound_flags & 0xC) >> 2
        self.sound_size = (sound_flags & 0x2) >> 1
        self.sound_type = sound_flags & 0x1

        if strict_parser():
            try:
                sound_format_to_string[self.sound_format]
            except KeyError:
                raise MalformedFLV("Invalid sound format: %d",
                                   self.sound_format)
            try:
                sound_rate_to_string[self.sound_rate]
            except KeyError:
                raise MalformedFLV("Invalid sound rate: %d",
                                   self.sound_rate)
            try:
                sound_size_to_string[self.sound_size]
            except KeyError:
                raise MalformedFLV("Invalid sound size type: %d",
                                   self.sound_size)
            try:
                sound_type_to_string[self.sound_type]
            except KeyError:
                raise MalformedFLV("Invalid sound type: %d",
                                   self.sound_type)

        if self.sound_format == SOUND_FORMAT_AAC:
            # AAC always has sampling rate of 44 kHz
            ensure(self.sound_rate, SOUND_RATE_44_KHZ,
                   "AAC sound format with incorrect sound rate: %d" %
                   self.sound_rate)
            # AAC is always stereo
            ensure(self.sound_type, SOUND_TYPE_STEREO,
                   "AAC sound format with incorrect sound type: %d" %
                   self.sound_type)

        f.seek(self.size - read_bytes, os.SEEK_CUR)

    def __repr__(self):
        if self.offset is None:
            return "<AudioTag unparsed>"
        elif self.sound_format is None:
            return ("<AudioTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<AudioTag at offset 0x%08X, time %d, size %d, %s>" %
                    (self.offset, self.timestamp, self.size,
                     sound_format_to_string[self.sound_format]))

class VideoTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.frame_type = None
        self.codec_id = None

    def parse_tag_content(self):
        f = self.f

        video_flags = get_ui8(f)
        read_bytes = 1

        self.frame_type = (video_flags & 0xF0) >> 4
        self.codec_id = video_flags & 0xF

        if strict_parser():
            try:
                frame_type_to_string[self.frame_type]
            except KeyError:
                raise MalformedFLV("Invalid frame type: %d", self.frame_type)
            try:
                codec_id_to_string[self.codec_id]
            except KeyError:
                raise MalformedFLV("Invalid codec ID: %d", self.codec_id)

        f.seek(self.size - read_bytes, os.SEEK_CUR)

    def __repr__(self):
        if self.offset is None:
            return "<VideoTag unparsed>"
        elif self.frame_type is None:
            return ("<VideoTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<VideoTag at offset 0x%08X, time %d, size %d, %s (%s)>" %
                    (self.offset, self.timestamp, self.size,
                     codec_id_to_string[self.codec_id],
                     frame_type_to_string[self.frame_type]))

class ScriptTag(Tag):

    def __init__(self, parent_flv, f):
        Tag.__init__(self, parent_flv, f)
        self.name = None
        self.variable = None

    def parse_tag_content(self):
        f = self.f

        # Here there's always a byte with the value of 0x02,
        # which means "string", although the spec says NOTHING
        # about it..
        value_type = get_ui8(f)
        ensure(value_type, 2, "The name of a script tag is not a string")

        # Need to pass the tag end offset, because apparently YouTube
        # doesn't give a *shit* about the FLV spec and just happily
        # ends the onMetaData tag after self.size bytes, instead of
        # ending it with the *required* 0x09 marker. Bastards!

        if strict_parser():
            # If we're strict, just don't pass this info
            tag_end = None
        else:
            # 11 = tag type (1) + data size (3) + timestamp (4) + stream id (3)
            tag_end = self.offset + 11 + self.size

        log.debug("max offset is 0x%08X", tag_end)

        self.name, self.variable = \
                   get_script_data_variable(f, max_offset=tag_end)
        log.debug("A script tag with a name of %s and value of %r",
                  self.name, self.variable)

    def __repr__(self):
        if self.offset is None:
            return "<ScriptTag unparsed>"
        elif not self.name:
            return ("<ScriptTag at offset 0x%08X, time %d, size %d>" %
                    (self.offset, self.timestamp, self.size))
        else:
            return ("<ScriptTag %s at offset 0x%08X, time %d, size %d>" %
                    (self.name, self.offset, self.timestamp, self.size))

tag_to_class = {
    TAG_TYPE_AUDIO: AudioTag,
    TAG_TYPE_VIDEO: VideoTag,
    TAG_TYPE_SCRIPT: ScriptTag
}

class FLV(object):

    def __init__(self, f):
        self.f = f
        self.version = None
        self.has_audio = None
        self.has_video = None
        self.tags = []

    def parse_header(self):
        f = self.f
        f.seek(0)

        # FLV header
        header = f.read(3)
        if len(header) < 3:
            raise MalformedFLV("The file is shorter than 3 bytes")

        # Do this irrelevant of STRICT_PARSING, to catch bogus files
        if header != "FLV":
            raise MalformedFLV("File signature is incorrect: 0x%X 0x%X 0x%X" %
                               struct.unpack("3B", header))

        # File version
        self.version = get_ui8(f)
        log.debug("File version is %d", self.version)

        # TypeFlags
        flags = get_ui8(f)

        ensure(flags & 0xF8, 0,
               "First TypeFlagsReserved field non zero: 0x%X" % (flags & 0xF8))
        ensure(flags & 0x2, 0,
               "Second TypeFlagsReserved field non zero: 0x%X" % (flags & 0x2))

        self.has_audio = False
        self.has_video = False
        if flags & 0x4:
            self.has_audio = True
        if flags & 0x1:
            self.has_video = True
        log.debug("File %s audio",
                  (self.has_audio and "has") or "does not have")
        log.debug("File %s video",
                  (self.has_video and "has") or "does not have")

        header_size = get_ui32(f)
        log.debug("Header size is %d bytes", header_size)

        f.seek(header_size)

        tag_0_size = get_ui32(f)
        ensure(tag_0_size, 0, "PreviousTagSize0 non zero: 0x%08X" % tag_0_size)

    def iter_tags(self):
        self.parse_header()
        try:
            while True:
                tag = self.get_next_tag()
                yield tag
        except EndOfTags:
            pass

    def read_tags(self):
        self.tags = list(self.iter_tags())

    def get_next_tag(self):
        f = self.f

        try:
            tag_type = get_ui8(f)
        except EndOfFile:
            raise EndOfTags

        tag_klass = self.tag_type_to_class(tag_type)
        tag = tag_klass(self, f)
        log.debug("Found a tag: %s", tag)

        tag.parse()

        return tag

    def tag_type_to_class(self, tag_type):
        try:
            return tag_to_class[tag_type]
        except KeyError:
            raise MalformedFLV("Invalid tag type: %d", tag_type)

def create_flv_tag(type, data, timestamp=0):
    tag_type = struct.pack("B", type)
    timestamp = make_si32_extended(timestamp)
    stream_id = make_ui24(0)

    data_size = len(data)
    tag_size = data_size + 11

    return ''.join([tag_type, make_ui24(data_size), timestamp, stream_id,
                    data, make_ui32(tag_size)])

def create_script_tag(name, data, timestamp=0):
    payload = make_ui8(2) + make_script_data_variable(name, data)
    return create_flv_tag(TAG_TYPE_SCRIPT, payload, timestamp)

def create_flv_header(has_audio=True, has_video=True):
    type_flags = 0
    if has_video:
        type_flags = type_flags | 0x1
    if has_audio:
        type_flags = type_flags | 0x4
    return ''.join(['FLV', make_ui8(1), make_ui8(type_flags), make_ui32(9),
                    make_ui32(0)])
<|endoftext|>


<commit_before>"""
Atmosphere API's extension of DRF permissions.
"""
from rest_framework import permissions

#NOTE: It would be better to use Object-level permissions here.
class ProjectOwnerRequired(permissions.BasePermission):
    def has_permission(self, request, view):
        auth_user = request.user
        project_id = view.kwargs.get('project_id')
        if not project_id:
            logger.warn("Could not find kwarg:'project_id'")
            return False
        return any(
                group for group in auth_user.group_set.all()
                if group.projects.filter(id=project_id))

class ApiAuthRequired(permissions.BasePermission):
    def has_permission(self, request, view):
        auth_user = request.user.is_authenticated()
        return auth_user

class ApiAuthOptional(permissions.BasePermission):
    def has_permission(self, request, view):
        #Allow access to GET/OPTIONS/HEAD operations.
        if request.method in permissions.SAFE_METHODS:
            return True
        #ASSERT: User trying to POST/PUT/PATCH, must be authenticated
        return request.user.is_authenticated()

class InMaintenance(permissions.BasePermission):
    def has_permission(self, request, view):
        return True
<commit_msg>add permission class for cloud admins<commit_after>"""
Atmosphere API's extension of DRF permissions.
"""
from rest_framework import permissions

from threepio import logger

# NOTE: It would be better to use Object-level permissions here.
class ProjectOwnerRequired(permissions.BasePermission):
    def has_permission(self, request, view):
        auth_user = request.user
        project_id = view.kwargs.get('project_id')
        if not project_id:
            logger.warn("Could not find kwarg:'project_id'")
            return False
        return any(
            group for group in auth_user.group_set.all()
            if group.projects.filter(id=project_id))

class ApiAuthRequired(permissions.BasePermission):
    def has_permission(self, request, view):
        auth_user = request.user.is_authenticated()
        return auth_user

class CloudAdminRequired(permissions.BasePermission):
    def has_permission(self, request, view):
        user = request.user

        if not user.is_authenticated():
            return False

        return user.is_staff or user.is_superuser

class ApiAuthOptional(permissions.BasePermission):
    def has_permission(self, request, view):
        # Allow access to GET/OPTIONS/HEAD operations.
        if request.method in permissions.SAFE_METHODS:
            return True
        # ASSERT: User trying to POST/PUT/PATCH, must be authenticated
        return request.user.is_authenticated()

class InMaintenance(permissions.BasePermission):
    def has_permission(self, request, view):
        return True
<|endoftext|>


<commit_before>"""
docstring
"""
<commit_msg>remove module docstring from python header<commit_after><|endoftext|>


<commit_before>import os, re;
# I got the actual size of the binary code wrong on the site once - this script should help prevent that.

dsDoc_by_sArch = {"w32": "x86", "w64": "x64", "win": "x86+x64"};
with open("build_info.txt", "rb") as oFile:
  iBuildNumber = int(re.search(r"build number\: (\d+)", oFile.read(), re.M).group(1));

print "Sizes (build %d)" % iBuildNumber;

for sArch in sorted(dsDoc_by_sArch.keys()):
  sDoc = dsDoc_by_sArch[sArch];
  iBinSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode.bin" % sArch);
  iBinESPSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-esp.bin" % sArch);
  try:
    iBinFuncSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-func.bin" % sArch);
    iBinCleanFuncSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-clean-func.bin" % sArch);
  except:
    print "  * %s: %d bytes (%d w/ stack allign)" % (sDoc, iBinSize, iBinESPSize);
  else:
    print "  * %s: %d bytes (%d w/ stack allign, %d w/ function wrapper, %d w/ function wrapper + save all registers)" % (sDoc, iBinSize, iBinESPSize, iBinFuncSize, iBinCleanFuncSize);
<commit_msg>Spelling: allign -> align<commit_after>import os, re;
# I got the actual size of the binary code wrong on the site once - this script should help prevent that.

dsDoc_by_sArch = {"w32": "x86", "w64": "x64", "win": "x86+x64"};
with open("build_info.txt", "rb") as oFile:
  iBuildNumber = int(re.search(r"build number\: (\d+)", oFile.read(), re.M).group(1));

print "Sizes (build %d)" % iBuildNumber;

for sArch in sorted(dsDoc_by_sArch.keys()):
  sDoc = dsDoc_by_sArch[sArch];
  iBinSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode.bin" % sArch);
  iBinESPSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-esp.bin" % sArch);
  try:
    iBinFuncSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-func.bin" % sArch);
    iBinCleanFuncSize = os.path.getsize(r"build\bin\%s-exec-calc-shellcode-clean-func.bin" % sArch);
  except:
    print "  * %s: %d bytes (%d w/ stack align)" % (sDoc, iBinSize, iBinESPSize);
  else:
    print "  * %s: %d bytes (%d w/ stack align, %d w/ function wrapper, %d w/ function wrapper + save all registers)" % (sDoc, iBinSize, iBinESPSize, iBinFuncSize, iBinCleanFuncSize);
<|endoftext|>