<commit_before>__author__ = 'vdowling'

# Standard imports
import json
import sys
import argparse

# 3rd party imports
import requests

# Local imports

class NaturalLanguageClassifierInstance(object):
    """
        Class that wraps the functionality associated with a single NLC instance. You can use this instead of command line POSTs to interact with the NLC
	There are three things that this make it very easy to do
		-Train a classifier
		-View information about what classifiers you have (and their statuses)
    """

    def __init__(self, username, password, url):
        if type(username) is not str or username == "":
            raise ValueError("Username provided is none or is not of proper type")
        elif type(password) is not str or password == "":
            raise ValueError("Password provided is none or is not of proper type")
        elif type(url) is not str or url == "":
            raise ValueError("Url provided is none or is not of proper type")
        else:
            self.username_ = username
            self.password_ = password
            self.url_ = url
            self.classifiers_ = list()
            resp = requests.get("%s/v1/classifiers" % self.url_, headers={"Content-Type": "application/json"}, auth=(self.username_, self.password_))
            if resp.ok:
                for classifier in resp.json().get('classifiers') if 'classifiers' in resp.json().keys() else []:
                    classifier_url = classifier.get('url')
                    classifier_obj = NaturalLanguageClassifier(self.username_, self.password_,
                                                               classifier_url=classifier_url)
                    self.classifiers_.append(classifier_obj)

    def train_classifier(self, classifier_name="CLASSIFIER", training_data=[], language='en', training_file=None):

        files = {'training_data': ('train-set.csv', open(training_file, 'rb'),'text/csv'), \
                 'training_metadata': (json.dumps({'language': 'en', 'name': classifier_name}))}
        resp = requests.post("%s/v1/classifiers" % self.url_,files=files, auth=(self.get_username(), self.get_password()))
        if resp.ok:
           classifier_url = resp.json().get('url')
           classifier_obj = NaturalLanguageClassifier(self.get_username(), self.get_password(), classifier_url=classifier_url)
           self.classifiers_.append(classifier_obj)
           print "Classifier created with ID " + classifier_obj.get_id()
           return classifier_obj
        else:
            raise resp.raise_for_status()

    def get_classifier_by_id(self, classifier_id=None):
        for cls in self.get_classifiers():
            if cls.get_id() == classifier_id:
                return cls
        return None

    def get_classifiers(self):
        return self.classifiers_

    def get_url(self):
        return self.url_

    def get_username(self):
        return self.username_

    def get_password(self):
        return self.password_

class NaturalLanguageClassifier(object):
    """
        Class that wraps the functionality for a single Natural Language Classifier
        This class assumes that the classifier has already been trained
    """

    def __init__(self, username, password, classifier_id=None, classifier_url=None):
        if type(username) is not str or username == "":
            raise ValueError("Username provided is none or is not of proper type")
        elif type(password) is not str or password == "":
            raise ValueError("Password provided is none or is not of proper type")
        elif classifier_id is None and classifier_url is None:
            raise ValueError("Both the classifier id and the url is None. At least one of these must be provided")
        elif classifier_url is not None:
            resp = requests.get(classifier_url, headers={"Content-Type": "application/json"}, auth=(username, password))
            if resp.ok:
                json_obj = resp.json()
                self.username_ = username
                self.password_ = password
                self.classifier_url_ = classifier_url
                self.classifier_id_ = json_obj.get('classifier_id')
                self.classifier_name_ = json_obj.get('name')
                self.created_ = json_obj.get('created')
            else:
                raise resp.raise_for_status()
        else:
            self.username_ = username
            self.password_ = password
            gateway_url = 'https://gateway.watsonplatform.net/natural-language-classifier/api/v1/classifiers/'
            resp = requests.get(gateway_url, headers={"Content-Type": "application/json"}, auth=(username, password))
            if resp.ok:
                classifier = [cls for cls in resp.json().get('classifiers') if
                              cls.get('classifier_id') == classifier_id]
                if len(classifier) == 0:
                    raise ValueError("No classifiers found with the id %r" % classifier_id)
                elif len(classifier) == 1:
                    cls = classifier[0]
                    self.classifier_url_ = cls.get('url')
                    self.classifier_name_ = cls.get('name')
                    self.created_ = cls.get('created')
                    self.classifier_id_ = classifier_id
            else:
                raise resp.raise_for_status()

    def get_id(self):
        return self.classifier_id_

    def get_name(self):
        return self.classifier_name_

    def get_created_date(self):
        return self.created_

    def get_url(self):
        return self.classifier_url_

    def get_status(self):
        resp = requests.get(self.get_url(), headers={"Content-Type": "application/json"},
                            auth=(self.username_, self.password_))
        if resp.ok:
            return resp.json().get('status')
        else:
            raise ValueError("Unable to get classifier status")

    def get_status_description(self):
        resp = requests.get(self.get_url(), headers={"Content-Type": "application/json"},
                            auth=(self.username_, self.password_))
        if resp.ok:
            return resp.json().get('status_description')
        else:
            raise ValueError("Unable to get classifier status")

    def classify(self, text):
        if type(text) is not str or text == "":
            raise ValueError("Provided text is not of proper type or is empty")
        else:
            json_data = {"text": text}
            resp = requests.post(self.get_url(), data=json.dumps(json_data),
                                 headers={"Content-Type": "application/json"}, auth=(self.username_, self.password_))
            if resp.ok:
                return resp.json()
            else:
                raise resp.raise_for_status()

#SAMPLE USAGE

    # 3. Get information on a specific classifier
    #classifierId = '{CLASSIFER_ID_FROM_STEP_1}'
    #print(' ')
   #print "Displaying classifier information for " + classifierId
   # nl_classifier = nlc_instance.get_classifier_by_id(classifierId)
    #print "Name : " + nl_classifier.get_name() + " Created on " + nl_classifier.get_created_date() + "Status: " + nlInstance.get_status()

    sys.exit(0)
<commit_msg>removed accidental sys.exit in nlc.py<commit_after>__author__ = 'vdowling'

# Standard imports
import json
import sys
import argparse

# 3rd party imports
import requests

# Local imports

class NaturalLanguageClassifierInstance(object):
    """
        Class that wraps the functionality associated with a single NLC instance. You can use this instead of command line POSTs to interact with the NLC
	There are three things that this make it very easy to do
		-Train a classifier
		-View information about what classifiers you have (and their statuses)
    """

    def __init__(self, username, password, url):
        if type(username) is not str or username == "":
            raise ValueError("Username provided is none or is not of proper type")
        elif type(password) is not str or password == "":
            raise ValueError("Password provided is none or is not of proper type")
        elif type(url) is not str or url == "":
            raise ValueError("Url provided is none or is not of proper type")
        else:
            self.username_ = username
            self.password_ = password
            self.url_ = url
            self.classifiers_ = list()
            resp = requests.get("%s/v1/classifiers" % self.url_, headers={"Content-Type": "application/json"}, auth=(self.username_, self.password_))
            if resp.ok:
                for classifier in resp.json().get('classifiers') if 'classifiers' in resp.json().keys() else []:
                    classifier_url = classifier.get('url')
                    classifier_obj = NaturalLanguageClassifier(self.username_, self.password_,
                                                               classifier_url=classifier_url)
                    self.classifiers_.append(classifier_obj)

    def train_classifier(self, classifier_name="CLASSIFIER", training_data=[], language='en', training_file=None):

        files = {'training_data': ('train-set.csv', open(training_file, 'rb'),'text/csv'), \
                 'training_metadata': (json.dumps({'language': 'en', 'name': classifier_name}))}
        resp = requests.post("%s/v1/classifiers" % self.url_,files=files, auth=(self.get_username(), self.get_password()))
        if resp.ok:
           classifier_url = resp.json().get('url')
           classifier_obj = NaturalLanguageClassifier(self.get_username(), self.get_password(), classifier_url=classifier_url)
           self.classifiers_.append(classifier_obj)
           print "Classifier created with ID " + classifier_obj.get_id()
           return classifier_obj
        else:
            raise resp.raise_for_status()

    def get_classifier_by_id(self, classifier_id=None):
        for cls in self.get_classifiers():
            if cls.get_id() == classifier_id:
                return cls
        return None

    def get_classifiers(self):
        return self.classifiers_

    def get_url(self):
        return self.url_

    def get_username(self):
        return self.username_

    def get_password(self):
        return self.password_

class NaturalLanguageClassifier(object):
    """
        Class that wraps the functionality for a single Natural Language Classifier
        This class assumes that the classifier has already been trained
    """

    def __init__(self, username, password, classifier_id=None, classifier_url=None):
        if type(username) is not str or username == "":
            raise ValueError("Username provided is none or is not of proper type")
        elif type(password) is not str or password == "":
            raise ValueError("Password provided is none or is not of proper type")
        elif classifier_id is None and classifier_url is None:
            raise ValueError("Both the classifier id and the url is None. At least one of these must be provided")
        elif classifier_url is not None:
            resp = requests.get(classifier_url, headers={"Content-Type": "application/json"}, auth=(username, password))
            if resp.ok:
                json_obj = resp.json()
                self.username_ = username
                self.password_ = password
                self.classifier_url_ = classifier_url
                self.classifier_id_ = json_obj.get('classifier_id')
                self.classifier_name_ = json_obj.get('name')
                self.created_ = json_obj.get('created')
            else:
                raise resp.raise_for_status()
        else:
            self.username_ = username
            self.password_ = password
            gateway_url = 'https://gateway.watsonplatform.net/natural-language-classifier/api/v1/classifiers/'
            resp = requests.get(gateway_url, headers={"Content-Type": "application/json"}, auth=(username, password))
            if resp.ok:
                classifier = [cls for cls in resp.json().get('classifiers') if
                              cls.get('classifier_id') == classifier_id]
                if len(classifier) == 0:
                    raise ValueError("No classifiers found with the id %r" % classifier_id)
                elif len(classifier) == 1:
                    cls = classifier[0]
                    self.classifier_url_ = cls.get('url')
                    self.classifier_name_ = cls.get('name')
                    self.created_ = cls.get('created')
                    self.classifier_id_ = classifier_id
            else:
                raise resp.raise_for_status()

    def get_id(self):
        return self.classifier_id_

    def get_name(self):
        return self.classifier_name_

    def get_created_date(self):
        return self.created_

    def get_url(self):
        return self.classifier_url_

    def get_status(self):
        resp = requests.get(self.get_url(), headers={"Content-Type": "application/json"},
                            auth=(self.username_, self.password_))
        if resp.ok:
            return resp.json().get('status')
        else:
            raise ValueError("Unable to get classifier status")

    def get_status_description(self):
        resp = requests.get(self.get_url(), headers={"Content-Type": "application/json"},
                            auth=(self.username_, self.password_))
        if resp.ok:
            return resp.json().get('status_description')
        else:
            raise ValueError("Unable to get classifier status")

    def classify(self, text):
        if type(text) is not str or text == "":
            raise ValueError("Provided text is not of proper type or is empty")
        else:
            json_data = {"text": text}
            resp = requests.post(self.get_url(), data=json.dumps(json_data),
                                 headers={"Content-Type": "application/json"}, auth=(self.username_, self.password_))
            if resp.ok:
                return resp.json()
            else:
                raise resp.raise_for_status()

#SAMPLE USAGE

    # 3. Get information on a specific classifier
    #classifierId = '{CLASSIFER_ID_FROM_STEP_1}'
    #print(' ')
   #print "Displaying classifier information for " + classifierId
   # nl_classifier = nlc_instance.get_classifier_by_id(classifierId)
    #print "Name : " + nl_classifier.get_name() + " Created on " + nl_classifier.get_created_date() + "Status: " + nlInstance.get_status()

<|endoftext|>


<commit_before>import csv
import re

from elasticsearch_dsl import analyzer, DocType, Object, Nested, String, Integer, Long

accent_analyzer = analyzer('accent_analyzer',
    tokenizer='standard',
    filter = ['lowercase', 'asciifolding']
)

class Constituent(DocType):
    ConstituentID = String()
    DisplayName = String(analyzer=accent_analyzer)
    DisplayDate = String()
    AlphaSort = String()
    Nationality = String()
    BeginDate = String()
    EndDate = String()
    ConstituentTypeID = String()
    addressTotal = Integer()
    nameSort = String()
    TextEntry = String()

    address = Nested(
        include_in_parent=True,
        properties={
            'ConAddressID' : String(),
            'ConstituentID' : String(),
            'AddressTypeID' : String(),
            'DisplayName2' : String(),
            'StreetLine1' : String(),
            'StreetLine2' : String(),
            'StreetLine3' : String(),
            'City' : String(),
            'State' : String(),
            'CountryID' : String(),
            'BeginDate' : String(),
            'EndDate' : String(),
            'Remarks' : String()
        }
    )

    biography = Object(
        properties={
            'TermID' : String(),
            'URL' : String()
        }
    )

    collection = Object(
        properties={
            'TermID' : String(),
            'URL' : String()
        }
    )

    format = Object(
        properties={
            'TermID' : String()
        }
    )

    gender = Object(
        properties={
            'TermID' : String()
        }
    )

    process = Object(
        properties={
            'TermID' : String()
        }
    )

    role = Object(
        properties={
            'TermID' : String()
        }
    )

class Converter:

    @staticmethod
    def remove_bom(row):
        cleaned = {}
        for k, v in row.iteritems():
            if (k.find('\xef\xbb\xbf')==0):
                kk = re.sub(r'\xef\xbb\xbf', r'', k)
                cleaned[kk] = v
            else:
                cleaned[k] = v
        return cleaned

    @staticmethod
    def convert_whitespace(string):
        string = string.replace('\xc2\xa0', ' ') # because WTF MSSQL Server!?
        string = string.replace('\u00a0', ' ') # because WTF MSSQL Server!?
        return string

    @staticmethod
    def remove_zeroes(string):
        while string.endswith("0") and len(string)>1:
            string = string[0:-1]
        return string

    @staticmethod
    def compress_address(remarks):
        precision = 4
        temp = remarks.replace(" ", "").split(",")
        lat = temp[0]
        if lat.find(".") != -1:
            lat = lat.split(".")
            lat[1] = Converter.remove_zeroes(lat[1][:precision])
            temp[0] = lat[0] + "." + lat[1]
        lon = temp[1]
        if lon.find(".") != -1:
            lon = lon.split(".")
            lon[1] = Converter.remove_zeroes(lon[1][:precision])
            temp[1] = lon[0] + "." + lon[1]
        return temp

    @staticmethod
    def process_csv(filename):
        print "loaded " + filename
        response = open(filename)
        reader = csv.DictReader(response)
        return reader

    @staticmethod
    def sort_addresses(addresses):
        sortedaddresses = []
        born = {}
        died = {}
        if len(addresses) <= 1:
            return addresses
        ## sort by date
        addresses = sorted(addresses, key=lambda d: d['BeginDate'])
        for add in addresses:
            # put the active/biz ones
            if (add['AddressTypeID'] == '7' or add['AddressTypeID'] == '2'):
                sortedaddresses.append(add)
            # find born if any
            if (add['AddressTypeID'] == '5'):
                born = add
            # find died if any
            if (add['AddressTypeID'] == '6'):
                died = add
        # prepend born
        if (born):
            sortedaddresses.insert(0,born)
        # append died
        if (died):
            sortedaddresses.append(died)
        return sortedaddresses
<commit_msg>:space_invader: fixed index sorting<commit_after>import csv
import re

from elasticsearch_dsl import analyzer, DocType, Object, Nested, String, Integer, Long

accent_analyzer = analyzer('accent_analyzer',
    tokenizer='standard',
    filter = ['lowercase', 'asciifolding']
)

class Constituent(DocType):
    ConstituentID = String()
    DisplayName = String(analyzer=accent_analyzer)
    DisplayDate = String()
    AlphaSort = String(
        analyzer=accent_analyzer,
        fields={'raw': String(index='not_analyzed')}
    )
    Nationality = String()
    BeginDate = String()
    EndDate = String()
    ConstituentTypeID = String()
    addressTotal = Integer()
    nameSort = String()
    TextEntry = String()

    address = Nested(
        include_in_parent=True,
        properties={
            'ConAddressID' : String(),
            'ConstituentID' : String(),
            'AddressTypeID' : String(),
            'DisplayName2' : String(),
            'StreetLine1' : String(),
            'StreetLine2' : String(),
            'StreetLine3' : String(),
            'City' : String(),
            'State' : String(),
            'CountryID' : String(),
            'BeginDate' : String(),
            'EndDate' : String(),
            'Remarks' : String()
        }
    )

    biography = Object(
        properties={
            'TermID' : String(),
            'URL' : String()
        }
    )

    collection = Object(
        properties={
            'TermID' : String(),
            'URL' : String()
        }
    )

    format = Object(
        properties={
            'TermID' : String()
        }
    )

    gender = Object(
        properties={
            'TermID' : String()
        }
    )

    process = Object(
        properties={
            'TermID' : String()
        }
    )

    role = Object(
        properties={
            'TermID' : String()
        }
    )

class Converter:

    @staticmethod
    def remove_bom(row):
        cleaned = {}
        for k, v in row.iteritems():
            if (k.find('\xef\xbb\xbf')==0):
                kk = re.sub(r'\xef\xbb\xbf', r'', k)
                cleaned[kk] = v
            else:
                cleaned[k] = v
        return cleaned

    @staticmethod
    def convert_whitespace(string):
        string = string.replace('\xc2\xa0', ' ') # because WTF MSSQL Server!?
        string = string.replace('\u00a0', ' ') # because WTF MSSQL Server!?
        return string

    @staticmethod
    def remove_zeroes(string):
        while string.endswith("0") and len(string)>1:
            string = string[0:-1]
        return string

    @staticmethod
    def compress_address(remarks):
        precision = 4
        temp = remarks.replace(" ", "").split(",")
        lat = temp[0]
        if lat.find(".") != -1:
            lat = lat.split(".")
            lat[1] = Converter.remove_zeroes(lat[1][:precision])
            temp[0] = lat[0] + "." + lat[1]
        lon = temp[1]
        if lon.find(".") != -1:
            lon = lon.split(".")
            lon[1] = Converter.remove_zeroes(lon[1][:precision])
            temp[1] = lon[0] + "." + lon[1]
        return temp

    @staticmethod
    def process_csv(filename):
        print "loaded " + filename
        response = open(filename)
        reader = csv.DictReader(response)
        return reader

    @staticmethod
    def sort_addresses(addresses):
        sortedaddresses = []
        born = {}
        died = {}
        if len(addresses) <= 1:
            return addresses
        ## sort by date
        addresses = sorted(addresses, key=lambda d: d['BeginDate'])
        for add in addresses:
            # put the active/biz ones
            if (add['AddressTypeID'] == '7' or add['AddressTypeID'] == '2'):
                sortedaddresses.append(add)
            # find born if any
            if (add['AddressTypeID'] == '5'):
                born = add
            # find died if any
            if (add['AddressTypeID'] == '6'):
                died = add
        # prepend born
        if (born):
            sortedaddresses.insert(0,born)
        # append died
        if (died):
            sortedaddresses.append(died)
        return sortedaddresses
<|endoftext|>


<commit_before>from pyspark.sql.column import Column
from pyspark.sql.functions import col
import sys
import dis

if sys.version >= '3':
    basestring = unicode = str
    long = int
    from io import StringIO
else:
    from cStringIO import StringIO

def for_name(name):
    """Dynamically load a class by its name.

    Equivalent to Java's Class.forName
    """
    lastdot = name.rfind('.')
    if (lastdot == -1):
        return getattr(__import('__main__'), name)

    mod = __import__(name[:lastdot])
    for comp in name.split('.')[1:]:
        mod = getattr(mod, comp)
    return mod

def smv_copy_array(sc, *cols):
    """Copy Python list to appropriate Java array
    """
    if (len(cols) == 0):        # may need to pass the correct java type somehow
        return sc._gateway.new_array(sc._jvm.java.lang.String, 0)

    elem = cols[0]
    if (isinstance(elem, basestring)):
        jcols = sc._gateway.new_array(sc._jvm.java.lang.String, len(cols))
        for i in range(0, len(jcols)):
            jcols[i] = cols[i]
    elif (isinstance(elem, Column)):
        jcols = sc._gateway.new_array(sc._jvm.org.apache.spark.sql.Column, len(cols))
        for i in range(0, len(jcols)):
            jcols[i] = cols[i]._jc
    elif (isinstance(elem, list)): # a list of list
        # use Java List as the outermost container; an Array[Array]
        # will not always work, because the inner list may be of
        # different lengths
        jcols = sc._jvm.java.util.ArrayList()
        for i in range(0, len(cols)):
            jcols.append(smv_copy_array(sc, *cols[i]))
    else:
        raise RuntimeError("Cannot copy array of type", type(elem))

    return jcols

def disassemble(mod):
    """Disassembles a module and returns bytecode as a string.
    """
    buf = StringIO()
    import dis
    if sys.version >= '3':
        dis.dis(mod, file=buf)
    else:
        stdout = sys.stdout
        sys.stdout = buf
        dis.dis(mod)
        sys.stdout = stdout
    ret = buf.getvalue()
    buf.close()
    return ret

# enhances the spark DataFrame with smv helper functions
from pyspark.sql import DataFrame

# peek from Scala side which print to STDOUT will not work on Jupyter. Have to pass the string to python side then print to stdout
import sys
DataFrame.peek = lambda df, pos = 1, colRegex = ".*": sys.stdout.write(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.peekStr(df._jdf, pos, colRegex) + "\n")

# provides df.selectPlus(...) in python shell
# example: df.selectPlus(lit(1).alias('col'))
# TODO: remove
DataFrame.selectPlus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.selectPlus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvExpandStruct = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvExpandStruct(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvGroupBy = lambda df, *cols: SmvGroupedData(df, df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvGroupBy(df._jdf, smv_copy_array(df._sc, *cols)))

DataFrame.smvJoinByKey = lambda df, other, keys, joinType: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvJoinByKey(df._jdf, other._jdf, smv_copy_array(df._sc, *keys), joinType), df.sql_ctx)

def __smvHashSample(df, key, rate=0.01, seed=23):
    if (isinstance(key, basestring)):
        jkey = col(key)._jc
    elif (isinstance(key, Column)):
        jkey = key._jc
    else:
        raise RuntimeError("key parameter must be either a String or a Column")
    return DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvHashSample(df._jdf, jkey, rate, seed), df.sql_ctx)
DataFrame.smvHashSample = __smvHashSample

DataFrame.smvSelectMinus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvSelectMinus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvSelectPlus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvSelectPlus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvDedupByKey = lambda df, keys: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvDedupByKey(df._jdf, smv_copy_array(df._sc, *keys)), df.sql_ctx)

import abc
from pyspark import SparkContext
from pyspark.sql import HiveContext

class Smv(object):
    """Creates a proxy to SmvApp.

    The SmvApp instance is exposed through the `app` attribute.
    """

    def init(self, arglist, _sc = None, _sqlContext = None):

        #TODO: appName should be read from the config files
        #      need to process the arglist first and create smvConfig before init SmvApp
        sc = SparkContext(appName="smvapp.py") if _sc is None else _sc
        sqlContext = HiveContext(sc) if _sqlContext is None else _sqlContext

        self.sqlContext = sqlContext
        self._jvm = sc._jvm

        # convert python arglist to java String array
        java_args =  smv_copy_array(sc, *arglist)
        self.app = self._jvm.org.tresamigos.smv.python.SmvPythonAppFactory.init(java_args, sqlContext._ssql_ctx)
        self.pymods = {}
        return self

    def runModule(self, fqn):
        """Runs a Scala SmvModule by its Fully Qualified Name(fqn)
        """
        jdf = self.app.app.runModuleByName(fqn)
        return DataFrame(jdf, self.sqlContext)

    def runDynamic(self, fqn):
        """Dynamically runs a Scala SmvModule by its Fully Qualified Name(fqn)
        """
        jdf = self.app.app.runDynamicModuleByName(fqn)
        return DataFrame(jdf, self.sqlContext)

    def run_python_module(self, name):
        klass = for_name(name)
        if (klass in self.pymods):
            return self.pymods[klass]
        else:
            return self.__resolve(klass, [klass])

    def createDF(self, schema, data):
        return DataFrame(self.app.dfFrom(schema, data), self.sqlContext)

    def __resolve(self, klass, stack):
        mod = klass(self)
        for dep in mod.requiresDS():
            if (dep in stack):
                raise RuntimeError("Circular module dependency detected", dep.fqn(), stack)

            stack.append(dep)
            res = self.__resolve(dep, stack)
            self.pymods[dep] = res
            stack.pop()

        tryRead = self.app.tryReadPersistedFile(mod.modulePath())
        if (tryRead.isSuccess()):
            ret = DataFrame(tryRead.get(), self.sqlContext)
        else:
            ret = mod.run(self.pymods)
            if not mod.isInput():
                self.app.persist(ret._jdf, mod.modulePath(), True)

        self.pymods[mod] = ret
        return ret

class SmvPyDataSet(object):
    """Base class for all SmvDataSets written in Python
    """

    __metaclass__ = abc.ABCMeta

    def __init__(self, smv):
        self.smv = smv

    @abc.abstractmethod
    def description(self):
        """A brief description of this dataset"""

    @abc.abstractmethod
    def requiresDS(self):
        """The list of dataset dependencies"""

    @abc.abstractmethod
    def run(self, i):
        """Comput this dataset, including its depencies if necessary"""

    def version(self):
        """All datasets are versioned, with a string,
        so that code and the data it produces can be tracked together."""
        return "0";

    def modulePath(self):
        return self.smv.app.outputDir() + "/" + self.fqn() + "_" + self.version() + ".csv"

    def fqn(self):
        """Returns the fully qualified name
        """
        return self.__module__ + "." + self.__class__.__name__

    def isInput(self):
        return False

class SmvPyCsvFile(SmvPyDataSet):
    """Raw input file in CSV format
    """

    def __init__(self, smv):
        super(SmvPyCsvFile, self).__init__(smv)
        self._smvCsvFile = smv.app.smvCsvFile(self.fqn() + "_" + self.version(), self.path(), self.csvAttr())

    def description(self):
        return "Input file: @" + self.path()

    @abc.abstractproperty
    def path(self):
        """The path to the csv input file"""

    def __mkCsvAttr(self, delimiter=',', quotechar='""', hasHeader=False):
        """Factory method for creating instances of Scala case class CsvAttributes"""
        return self.smv._jvm.org.tresamigos.smv.CsvAttributes(delimiter, quotechar, hasHeader)

    def defaultCsvWithHeader(self):
        return self.__mkCsvAttr(hasHeader=True)

    def defaultTsv(self):
        return self.__mkCsvAttr(delimiter='\t')

    def defaultTsvWitHeader(self):
        return self.__mkCsvAttr(delimier='\t', hasHeader=True)

    def csvAttr(self):
        """Specifies the csv file format.  Corresponds to the CsvAttributes case class in Scala.
        """
        return None

    def isInput(self):
        return True

    def requiresDS(self):
        return []

    def run(self, i):
        jdf = self._smvCsvFile.rdd()
        return DataFrame(jdf, self.smv.sqlContext)

class SmvPyHiveTable(SmvPyDataSet):
    """Input data source from a Hive table
    """

    def __init__(self, smv):
        super(SmvPyHiveTable, self).__init__(smv)
        self._smvHiveTable = self.smv._jvm.org.tresamigos.smv.SmvHiveTable(self.tableName())

    def description(self):
        return "Hive Table: @" + self.tableName()

    @abc.abstractproperty
    def tableName(self):
        """The qualified Hive table name"""

    def isInput(self):
        return True

    def requiresDS(self):
        return []

    def run(self, i):
        return DataFrame(self._smvHiveTable.rdd(), self.smv.sqlContext)

class SmvPyModule(SmvPyDataSet):
    """Base class for SmvModules written in Python
    """

    def __init__(self, smv):
        super(SmvPyModule, self).__init__(smv)

    def prerun(self, i):
        print(".... computing module " + self.fqn())

class SmvGroupedData(object):
    """Wrapper around the Scala SmvGroupedData"""
    def __init__(self, df, sgd):
        self.df = df
        self.sgd = sgd

    def smvTopNRecs(self, maxElems, *cols):
        return DataFrame(self.sgd.smvTopNRecs(maxElems, smv_copy_array(self.df._sc, *cols)), self.df.sql_ctx)

    def smvPivotSum(self, pivotCols, valueCols, baseOutput):
        """Perform a normal SmvPivot using the pivot and value columns, followed by summing all the output.
        TODO: docstring parameters and return values
        pivotCols: list of list of strings
        valueCols: list of strings
        baseOutput: list of strings"""
        return DataFrame(self.sgd.smvPivotSum(smv_copy_array(self.df._sc, *pivotCols), smv_copy_array(self.df._sc, *valueCols), smv_copy_array(self.df._sc, *baseOutput)), self.df.sql_ctx)

# Create the SmvApp "Singleton"
SmvApp = Smv()

# ColumnHelper methods:
Column.smvStrToTimestamp = lambda c, fmt: Column(SmvApp._jvm.org.tresamigos.smv.ColumnHelper(c._jc).smvStrToTimestamp(fmt))
<commit_msg>feat(#383: cut 1 for Python module auto-version)<commit_after>from pyspark.sql.column import Column
from pyspark.sql.functions import col
import sys

if sys.version >= '3':
    basestring = unicode = str
    long = int
    from io import StringIO
else:
    from cStringIO import StringIO

def for_name(name):
    """Dynamically load a class by its name.

    Equivalent to Java's Class.forName
    """
    lastdot = name.rfind('.')
    if (lastdot == -1):
        return getattr(__import('__main__'), name)

    mod = __import__(name[:lastdot])
    for comp in name.split('.')[1:]:
        mod = getattr(mod, comp)
    return mod

def smv_copy_array(sc, *cols):
    """Copy Python list to appropriate Java array
    """
    if (len(cols) == 0):        # may need to pass the correct java type somehow
        return sc._gateway.new_array(sc._jvm.java.lang.String, 0)

    elem = cols[0]
    if (isinstance(elem, basestring)):
        jcols = sc._gateway.new_array(sc._jvm.java.lang.String, len(cols))
        for i in range(0, len(jcols)):
            jcols[i] = cols[i]
    elif (isinstance(elem, Column)):
        jcols = sc._gateway.new_array(sc._jvm.org.apache.spark.sql.Column, len(cols))
        for i in range(0, len(jcols)):
            jcols[i] = cols[i]._jc
    elif (isinstance(elem, list)): # a list of list
        # use Java List as the outermost container; an Array[Array]
        # will not always work, because the inner list may be of
        # different lengths
        jcols = sc._jvm.java.util.ArrayList()
        for i in range(0, len(cols)):
            jcols.append(smv_copy_array(sc, *cols[i]))
    else:
        raise RuntimeError("Cannot copy array of type", type(elem))

    return jcols

def disassemble(obj):
    """Disassembles a module and returns bytecode as a string.
    """
    mod = obj if (isinstance(obj, type)) else obj.__class__

    buf = StringIO()
    import dis
    if sys.version >= '3':
        dis.dis(mod, file=buf)
    else:
        stdout = sys.stdout
        sys.stdout = buf
        dis.dis(mod)
        sys.stdout = stdout
    ret = buf.getvalue()
    buf.close()
    return ret

# enhances the spark DataFrame with smv helper functions
from pyspark.sql import DataFrame

# peek from Scala side which print to STDOUT will not work on Jupyter. Have to pass the string to python side then print to stdout
import sys
DataFrame.peek = lambda df, pos = 1, colRegex = ".*": sys.stdout.write(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.peekStr(df._jdf, pos, colRegex) + "\n")

# provides df.selectPlus(...) in python shell
# example: df.selectPlus(lit(1).alias('col'))
# TODO: remove
DataFrame.selectPlus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.selectPlus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvExpandStruct = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvExpandStruct(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvGroupBy = lambda df, *cols: SmvGroupedData(df, df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvGroupBy(df._jdf, smv_copy_array(df._sc, *cols)))

DataFrame.smvJoinByKey = lambda df, other, keys, joinType: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvJoinByKey(df._jdf, other._jdf, smv_copy_array(df._sc, *keys), joinType), df.sql_ctx)

def __smvHashSample(df, key, rate=0.01, seed=23):
    if (isinstance(key, basestring)):
        jkey = col(key)._jc
    elif (isinstance(key, Column)):
        jkey = key._jc
    else:
        raise RuntimeError("key parameter must be either a String or a Column")
    return DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvHashSample(df._jdf, jkey, rate, seed), df.sql_ctx)
DataFrame.smvHashSample = __smvHashSample

DataFrame.smvSelectMinus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvSelectMinus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvSelectPlus = lambda df, *cols: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvSelectPlus(df._jdf, smv_copy_array(df._sc, *cols)), df.sql_ctx)

DataFrame.smvDedupByKey = lambda df, keys: DataFrame(df._sc._jvm.org.tresamigos.smv.python.SmvPythonHelper.smvDedupByKey(df._jdf, smv_copy_array(df._sc, *keys)), df.sql_ctx)

import abc
from pyspark import SparkContext
from pyspark.sql import HiveContext

class Smv(object):
    """Creates a proxy to SmvApp.

    The SmvApp instance is exposed through the `app` attribute.
    """

    def init(self, arglist, _sc = None, _sqlContext = None):

        #TODO: appName should be read from the config files
        #      need to process the arglist first and create smvConfig before init SmvApp
        sc = SparkContext(appName="smvapp.py") if _sc is None else _sc
        sqlContext = HiveContext(sc) if _sqlContext is None else _sqlContext

        self.sqlContext = sqlContext
        self._jvm = sc._jvm

        # convert python arglist to java String array
        java_args =  smv_copy_array(sc, *arglist)
        self.app = self._jvm.org.tresamigos.smv.python.SmvPythonAppFactory.init(java_args, sqlContext._ssql_ctx)
        self.pymods = {}
        return self

    def runModule(self, fqn):
        """Runs a Scala SmvModule by its Fully Qualified Name(fqn)
        """
        jdf = self.app.app.runModuleByName(fqn)
        return DataFrame(jdf, self.sqlContext)

    def runDynamic(self, fqn):
        """Dynamically runs a Scala SmvModule by its Fully Qualified Name(fqn)
        """
        jdf = self.app.app.runDynamicModuleByName(fqn)
        return DataFrame(jdf, self.sqlContext)

    def run_python_module(self, name):
        klass = for_name(name)
        if (klass in self.pymods):
            return self.pymods[klass]
        else:
            return self.__resolve(klass, [klass])

    def createDF(self, schema, data):
        return DataFrame(self.app.dfFrom(schema, data), self.sqlContext)

    def __resolve(self, klass, stack):
        mod = klass(self)
        for dep in mod.requiresDS():
            if (dep in stack):
                raise RuntimeError("Circular module dependency detected", dep.fqn(), stack)

            stack.append(dep)
            res = self.__resolve(dep, stack)
            self.pymods[dep] = res
            stack.pop()

        tryRead = self.app.tryReadPersistedFile(mod.modulePath())
        if (tryRead.isSuccess()):
            ret = DataFrame(tryRead.get(), self.sqlContext)
        else:
            ret = mod.run(self.pymods)
            if not mod.isInput():
                self.app.persist(ret._jdf, mod.modulePath(), True)

        self.pymods[mod] = ret
        return ret

class SmvPyDataSet(object):
    """Base class for all SmvDataSets written in Python
    """

    __metaclass__ = abc.ABCMeta

    def __init__(self, smv):
        self.smv = smv

    @abc.abstractmethod
    def description(self):
        """A brief description of this dataset"""

    @abc.abstractmethod
    def requiresDS(self):
        """The list of dataset dependencies"""

    @abc.abstractmethod
    def run(self, i):
        """Comput this dataset, including its depencies if necessary"""

    def version(self):
        """All datasets are versioned, with a string,
        so that code and the data it produces can be tracked together."""
        return "0";

    @classmethod
    def datasetHash(cls):
        try:
            import inspect
            src = inspect.getsource(cls)
            code = compile(src, inspect.getsourcefile(cls), 'exec')
            res = hash(code)
        except: # `inspect` will raise error for classes defined in the REPL
            res = hash(disassemble(cls))
        return res

    def hashOfHash(self):
        res = hash(self.version() + str(self.datasetHash()))

        # include datasetHash of dependency modules
        if sys.version >= '3':
            from functools import reduce
        res = reduce(lambda x,y: x+y, [m.datasetHash() for m in self.requiresDS()], res)

        return res

    def modulePath(self):
        return self.smv.app.outputDir() + "/" + self.fqn() + "_" + hex(self.hashOfHash() & 0xffffffff)[2:] + ".csv"

    def fqn(self):
        """Returns the fully qualified name
        """
        return self.__module__ + "." + self.__class__.__name__

    def isInput(self):
        return False

class SmvPyCsvFile(SmvPyDataSet):
    """Raw input file in CSV format
    """

    def __init__(self, smv):
        super(SmvPyCsvFile, self).__init__(smv)
        self._smvCsvFile = smv.app.smvCsvFile(self.fqn() + "_" + self.version(), self.path(), self.csvAttr())

    def description(self):
        return "Input file: @" + self.path()

    @abc.abstractproperty
    def path(self):
        """The path to the csv input file"""

    def __mkCsvAttr(self, delimiter=',', quotechar='""', hasHeader=False):
        """Factory method for creating instances of Scala case class CsvAttributes"""
        return self.smv._jvm.org.tresamigos.smv.CsvAttributes(delimiter, quotechar, hasHeader)

    def defaultCsvWithHeader(self):
        return self.__mkCsvAttr(hasHeader=True)

    def defaultTsv(self):
        return self.__mkCsvAttr(delimiter='\t')

    def defaultTsvWitHeader(self):
        return self.__mkCsvAttr(delimier='\t', hasHeader=True)

    def csvAttr(self):
        """Specifies the csv file format.  Corresponds to the CsvAttributes case class in Scala.
        """
        return None

    def isInput(self):
        return True

    def requiresDS(self):
        return []

    def run(self, i):
        jdf = self._smvCsvFile.rdd()
        return DataFrame(jdf, self.smv.sqlContext)

class SmvPyHiveTable(SmvPyDataSet):
    """Input data source from a Hive table
    """

    def __init__(self, smv):
        super(SmvPyHiveTable, self).__init__(smv)
        self._smvHiveTable = self.smv._jvm.org.tresamigos.smv.SmvHiveTable(self.tableName())

    def description(self):
        return "Hive Table: @" + self.tableName()

    @abc.abstractproperty
    def tableName(self):
        """The qualified Hive table name"""

    def isInput(self):
        return True

    def requiresDS(self):
        return []

    def run(self, i):
        return DataFrame(self._smvHiveTable.rdd(), self.smv.sqlContext)

class SmvPyModule(SmvPyDataSet):
    """Base class for SmvModules written in Python
    """

    def __init__(self, smv):
        super(SmvPyModule, self).__init__(smv)

    def prerun(self, i):
        print(".... computing module " + self.fqn())

class SmvGroupedData(object):
    """Wrapper around the Scala SmvGroupedData"""
    def __init__(self, df, sgd):
        self.df = df
        self.sgd = sgd

    def smvTopNRecs(self, maxElems, *cols):
        return DataFrame(self.sgd.smvTopNRecs(maxElems, smv_copy_array(self.df._sc, *cols)), self.df.sql_ctx)

    def smvPivotSum(self, pivotCols, valueCols, baseOutput):
        """Perform a normal SmvPivot using the pivot and value columns, followed by summing all the output.
        TODO: docstring parameters and return values
        pivotCols: list of list of strings
        valueCols: list of strings
        baseOutput: list of strings"""
        return DataFrame(self.sgd.smvPivotSum(smv_copy_array(self.df._sc, *pivotCols), smv_copy_array(self.df._sc, *valueCols), smv_copy_array(self.df._sc, *baseOutput)), self.df.sql_ctx)

# Create the SmvApp "Singleton"
SmvApp = Smv()

# ColumnHelper methods:
Column.smvStrToTimestamp = lambda c, fmt: Column(SmvApp._jvm.org.tresamigos.smv.ColumnHelper(c._jc).smvStrToTimestamp(fmt))
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import json
import logging
import re
import warnings
try:
    from urllib2 import urlopen
    from urlparse import urlparse, parse_qs, unquote
except ImportError:
    from urllib.parse import urlparse, parse_qs, unquote
    from urllib.request import urlopen

from .exceptions import MultipleObjectsReturned, PytubeError, CipherError, \
    DoesNotExist, AgeRestricted
from .jsinterp import JSInterpreter
from .models import Video
from .utils import safe_filename

log = logging.getLogger(__name__)

# YouTube quality and codecs id map.
YT_QUALITY_PROFILES = {
    # flash
    5: ["flv", "240p", "Sorenson H.263", "N/A", "0.25", "MP3", "64"],

    # 3gp
    17: ["3gp", "144p", "MPEG-4 Visual", "Simple", "0.05", "AAC", "24"],
    36: ["3gp", "240p", "MPEG-4 Visual", "Simple", "0.17", "AAC", "38"],

    # webm
    43: ["webm", "360p", "VP8", "N/A", "0.5", "Vorbis", "128"],
    100: ["webm", "360p", "VP8", "3D", "N/A", "Vorbis", "128"],

    # mpeg4
    18: ["mp4", "360p", "H.264", "Baseline", "0.5", "AAC", "96"],
    22: ["mp4", "720p", "H.264", "High", "2-2.9", "AAC", "192"],
    82: ["mp4", "360p", "H.264", "3D", "0.5", "AAC", "96"],
    83: ["mp4", "240p", "H.264", "3D", "0.5", "AAC", "96"],
    84: ["mp4", "720p", "H.264", "3D", "2-2.9", "AAC", "152"],
    85: ["mp4", "1080p", "H.264", "3D", "2-2.9", "AAC", "152"],
}

# The keys corresponding to the quality/codec map above.
YT_QUALITY_PROFILE_KEYS = (
    'extension',
    'resolution',
    'video_codec',
    'profile',
    'video_bitrate',
    'audio_codec',
    'audio_bitrate'
)

class YouTube(object):
    """Class representation of a single instance of a YouTube session.
    """
    def __init__(self, url=None):
        """Initializes YouTube API wrapper.

        :param str url:
            The url to the YouTube video.
        """
        self._filename = None
        self._video_url = None
        self._js_code = False
        self._videos = []
        if url:
            self.from_url(url)

    @property
    def url(self):
        """Gets the video url."""
        return self._video_url

    @url.setter
    def url(self, url):
        """Sets the url for the video (This method is deprecated. Use
        ``from_url()`` instead).

        :param str url:
            The url to the YouTube video.
        """
        warnings.warn("url setter deprecated, use ``from_url()`` "
                      "instead.", DeprecationWarning)
        self.from_url(url)

    @property
    def video_id(self):
        """Gets the video id by parsing and extracting it from the url."""
        parts = urlparse(self._video_url)
        qs = getattr(parts, 'query')
        if qs:
            video_id = parse_qs(qs).get('v')
            if video_id:
                return video_id.pop()
        return False

    @property
    def filename(self):
        """Gets the filename of the video.  If it hasn't been defined by the
        user, the title will instead be used.
        """
        if not self._filename:
            self._filename = safe_filename(self.title)
            log.debug("generated 'safe' filename: %s", self._filename)
        return self._filename

    @filename.setter
    def filename(self, filename):
        """Sets the filename (This method is deprecated. Use ``set_filename()``
        instead).

        :param str filename:
            The filename of the video.
        """
        warnings.warn("filename setter deprecated. Use ``set_filename()`` "
                      "instead.", DeprecationWarning)
        self.set_filename(filename)

    def set_filename(self, filename):
        """Sets the filename of the video.

        :param str filename:
            The filename of the video.
        """
        # TODO: Check if the filename contains the file extension and either
        # strip it or raise an exception.
        self._filename = filename
        if self.get_videos():
            for video in self.get_videos():
                video.filename = filename
        return True

    def get_videos(self):
        """Gets all videos."""
        return self._videos

    @property
    def videos(self):
        """Gets all videos. (This method is deprecated. Use ``get_videos()``
        instead.
        """
        warnings.warn("videos property deprecated. Use ``get_videos()`` "
                      "instead.", DeprecationWarning)
        return self._videos

    def from_url(self, url):
        """Sets the url for the video.

        :param str url:
            The url to the YouTube video.
        """
        self._video_url = url

        # Reset the filename incase it was previously set.
        self._filename = None

        # Get the video details.
        video_data = self.get_video_data()

        # Set the title from the title.
        self.title = video_data.get("args", {}).get("title")

        # Rewrite and add the url to the javascript file, we'll need to fetch
        # this if YouTube doesn't provide us with the signature.
        js_url = "http:" + video_data.get("assets", {}).get("js")

        # Just make these easily accessible as variables.
        stream_map = video_data.get("args", {}).get("stream_map")
        video_urls = stream_map.get("url")

        # For each video url, identify the quality profile and add it to list
        # of available videos.
        for idx, url in enumerate(video_urls):
            log.debug("attempting to get quality profile from url: %s", url)
            try:
                itag, quality_profile = self._get_quality_profile_from_url(url)
                if not quality_profile:
                    log.warn("unable to identify profile for itag=%s", itag)
                    continue
            except (TypeError, KeyError) as e:
                log.exception("passing on exception %s", e)
                continue

            # Check if we have the signature, otherwise we'll need to get the
            # cipher from the js.
            if "signature=" not in url:
                log.debug('signature not in url, attempting to resolve the '
                          'cipher...')
                signature = self._get_cipher(stream_map["s"][idx], js_url)
                url = "{}&signature={}".format(url, signature)
            self._add_video(url, self.filename, **quality_profile)

    def get(self, extension=None, resolution=None, profile=None):
        """Gets a single video given a file extention (and/or resolution
        and/or quality profile).

        :param str extention:
            The desired file extention (e.g.: mp4, flv).
        :param str resolution:
            The desired video broadcasting standard (e.g.: 720p, 1080p)
        :param str profile:
            The desired quality profile (this is subjective, I don't recommend
            using it).
        """
        result = []
        for v in self.get_videos():
            if extension and v.extension != extension:
                continue
            elif resolution and v.resolution != resolution:
                continue
            elif profile and v.profile != profile:
                continue
            else:
                result.append(v)
        matches = len(result)
        if matches <= 0:
            return DoesNotExist("No videos met this criteria.")
        elif matches == 1:
            return result[0]
        else:
            raise MultipleObjectsReturned("Multiple videos met this criteria.")

    def filter(self, extension=None, resolution=None, profile=None):
        """Gets a filtered list of videos given a file extention and/or
        resolution and/or quality profile.

        :param str extention:
            The desired file extention (e.g.: mp4, flv).
        :param str resolution:
            The desired video broadcasting standard (e.g.: 720p, 1080p)
        :param str profile:
            The desired quality profile (this is subjective, I don't recommend
            using it).
        """
        results = []
        for v in self.get_videos():
            if extension and v.extension != extension:
                continue
            elif resolution and v.resolution != resolution:
                continue
            elif profile and v.profile != profile:
                continue
            else:
                results.append(v)
        return results

    def get_video_data(self):
        """Gets the page and extracts out the video data."""
        # Reset the filename incase it was previously set.
        self.title = None
        response = urlopen(self.url)
        if not response:
            raise PytubeError("Unable to open url: {}".format(self.url))

        html = response.read().decode("utf-8")
        if "og:restrictions:age" in html:
            raise AgeRestricted("Age restricted video. Unable to download "
                                "without being signed in.")

        # Extract out the json data from the html response body.
        json_object = self._get_json_data(html)

        # Here we decode the stream map and bundle it into the json object. We
        # do this just so we just can return one object for the video data.
        encoded_stream_map = json_object.get("args", {}).get(
            "url_encoded_fmt_stream_map")
        json_object['args']['stream_map'] = self._parse_stream_map(
            encoded_stream_map)
        return json_object

    def _parse_stream_map(self, blob):
        """A modified version of ``urlparse.parse_qs`` that's able to decode
        YouTube's stream map.

        :param str blob:
            An encoded blob of text containing the stream map data.
        """
        dct = {
            "itag": [],
            "url": [],
            "quality": [],
            "fallback_host": [],
            "s": [],
            "type": []
        }

        # Split the comma separated videos.
        videos = blob.split(",")

        # Unquote the characters and split to parameters.
        videos = [video.split("&") for video in videos]

        # Split at the equals sign so we can break this key value pairs and
        # toss it into a dictionary.
        for video in videos:
            for kv in video:
                key, value = kv.split("=")
                dct.get(key, []).append(unquote(value))
        log.debug('decoded stream map: %s', dct)
        return dct

    def _get_json_data(self, html):
        """Extract the json out from the html.

        :param str html:
            The raw html of the page.
        """
        # 18 represents the length of "ytplayer.config = ".
        start = html.find("ytplayer.config = ") + 18
        html = html[start:]

        offset = self._get_json_offset(html)
        if not offset:
            raise PytubeError("Unable to extract json.")
        return json.loads(html[:offset])

    def _get_json_offset(self, html):
        """Find where the json object starts.

        :param str html:
            The raw html of the YouTube page.
        """
        brackets = []
        index = 1
        # Determine the offset by pushing/popping brackets until all
        # js expressions are closed.
        for idx, ch in enumerate(html):
            if ch == "{":
                brackets.append("}")
            elif ch == "}":
                brackets.pop()
                if len(brackets) == 0:
                    break
        else:
            raise PytubeError("Unable to determine json offset.")
        return index + idx

    def _get_cipher(self, signature, url):
        """Gets the signature using the cipher.

        :param str signature:
            The url signature.
        :param str url:
            The url of the javascript file.
        """
        reg_exp = re.compile(r'\.sig\|\|([a-zA-Z0-9$]+)\(')
        if not self._js_code:
            response = urlopen(url)
            if not response:
                raise PytubeError("Unable to open url: {}".format(self.url))
            self._js_code = response.read().decode("utf-8")
        try:
            matches = reg_exp.search(self._js_code)
            if matches:
                # Return the first matching group.
                func = next(g for g in matches.groups() if g is not None)
            # Load js into JS Python interpreter.
            jsi = JSInterpreter(self._js_code)
            initial_function = jsi.extract_function(func)
            return initial_function([signature])
        except Exception as e:
            raise CipherError("Couldn't cipher the signature. Maybe YouTube "
                              "has changed the cipher algorithm. Notify this "
                              "issue on GitHub: {}".format(e))
        return False

    def _get_quality_profile_from_url(self, video_url):
        """Gets the quality profile given a video url. Normally we would just
        use ``urlparse`` since itags are represented as a get parameter, but
        YouTube doesn't pass a properly encoded url.

        :param str video_url:
            The malformed url-encoded video_url.
        """
        reg_exp = re.compile('itag=(\d+)')
        itag = reg_exp.findall(video_url)
        if itag and len(itag) == 1:
            itag = int(itag[0])
            # Given an itag, refer to the YouTube quality profiles to get the
            # properties (media type, resolution, etc.) of the video.
            quality_profile = YT_QUALITY_PROFILES.get(itag)
            if not quality_profile:
                return itag, None
            # Here we just combine the quality profile keys to the
            # corresponding quality profile, referenced by the itag.
            return itag, dict(zip(YT_QUALITY_PROFILE_KEYS, quality_profile))
        if not itag:
            raise PytubeError("Unable to get encoding profile, no itag found.")
        elif len(itag) > 1:
            log.warn("Multiple itags found: %s", itag)
            raise PytubeError("Unable to get encoding profile, multiple itags "
                              "found.")
        return False

    def _add_video(self, url, filename, **kwargs):
        """Adds new video object to videos.

        :param str url:
            The signed url to the video.
        :param str filename:
            The filename for the video.
        :param kwargs:
            Additional properties to set for the video object.
        """
        video = Video(url, filename, **kwargs)
        self._videos.append(video)
        self._videos.sort()
        return True
<commit_msg>these should really be immutable<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import json
import logging
import re
import warnings
try:
    from urllib2 import urlopen
    from urlparse import urlparse, parse_qs, unquote
except ImportError:
    from urllib.parse import urlparse, parse_qs, unquote
    from urllib.request import urlopen

from .exceptions import MultipleObjectsReturned, PytubeError, CipherError, \
    DoesNotExist, AgeRestricted
from .jsinterp import JSInterpreter
from .models import Video
from .utils import safe_filename

log = logging.getLogger(__name__)

# YouTube quality and codecs id map.
YT_QUALITY_PROFILES = {
    # flash
    5: ("flv", "240p", "Sorenson H.263", "N/A", "0.25", "MP3", "64"),

    # 3gp
    17: ("3gp", "144p", "MPEG-4 Visual", "Simple", "0.05", "AAC", "24"),
    36: ("3gp", "240p", "MPEG-4 Visual", "Simple", "0.17", "AAC", "38"),

    # webm
    43: ("webm", "360p", "VP8", "N/A", "0.5", "Vorbis", "128"),
    100: ("webm", "360p", "VP8", "3D", "N/A", "Vorbis", "128"),

    # mpeg4
    18: ("mp4", "360p", "H.264", "Baseline", "0.5", "AAC", "96"),
    22: ("mp4", "720p", "H.264", "High", "2-2.9", "AAC", "192"),
    82: ("mp4", "360p", "H.264", "3D", "0.5", "AAC", "96"),
    83: ("mp4", "240p", "H.264", "3D", "0.5", "AAC", "96"),
    84: ("mp4", "720p", "H.264", "3D", "2-2.9", "AAC", "152"),
    85: ("mp4", "1080p", "H.264", "3D", "2-2.9", "AAC", "152"),
}

# The keys corresponding to the quality/codec map above.
YT_QUALITY_PROFILE_KEYS = (
    'extension',
    'resolution',
    'video_codec',
    'profile',
    'video_bitrate',
    'audio_codec',
    'audio_bitrate'
)

class YouTube(object):
    """Class representation of a single instance of a YouTube session.
    """
    def __init__(self, url=None):
        """Initializes YouTube API wrapper.

        :param str url:
            The url to the YouTube video.
        """
        self._filename = None
        self._video_url = None
        self._js_code = False
        self._videos = []
        if url:
            self.from_url(url)

    @property
    def url(self):
        """Gets the video url."""
        return self._video_url

    @url.setter
    def url(self, url):
        """Sets the url for the video (This method is deprecated. Use
        ``from_url()`` instead).

        :param str url:
            The url to the YouTube video.
        """
        warnings.warn("url setter deprecated, use ``from_url()`` "
                      "instead.", DeprecationWarning)
        self.from_url(url)

    @property
    def video_id(self):
        """Gets the video id by parsing and extracting it from the url."""
        parts = urlparse(self._video_url)
        qs = getattr(parts, 'query')
        if qs:
            video_id = parse_qs(qs).get('v')
            if video_id:
                return video_id.pop()
        return False

    @property
    def filename(self):
        """Gets the filename of the video.  If it hasn't been defined by the
        user, the title will instead be used.
        """
        if not self._filename:
            self._filename = safe_filename(self.title)
            log.debug("generated 'safe' filename: %s", self._filename)
        return self._filename

    @filename.setter
    def filename(self, filename):
        """Sets the filename (This method is deprecated. Use ``set_filename()``
        instead).

        :param str filename:
            The filename of the video.
        """
        warnings.warn("filename setter deprecated. Use ``set_filename()`` "
                      "instead.", DeprecationWarning)
        self.set_filename(filename)

    def set_filename(self, filename):
        """Sets the filename of the video.

        :param str filename:
            The filename of the video.
        """
        # TODO: Check if the filename contains the file extension and either
        # strip it or raise an exception.
        self._filename = filename
        if self.get_videos():
            for video in self.get_videos():
                video.filename = filename
        return True

    def get_videos(self):
        """Gets all videos."""
        return self._videos

    @property
    def videos(self):
        """Gets all videos. (This method is deprecated. Use ``get_videos()``
        instead.
        """
        warnings.warn("videos property deprecated. Use ``get_videos()`` "
                      "instead.", DeprecationWarning)
        return self._videos

    def from_url(self, url):
        """Sets the url for the video.

        :param str url:
            The url to the YouTube video.
        """
        self._video_url = url

        # Reset the filename incase it was previously set.
        self._filename = None

        # Get the video details.
        video_data = self.get_video_data()

        # Set the title from the title.
        self.title = video_data.get("args", {}).get("title")

        # Rewrite and add the url to the javascript file, we'll need to fetch
        # this if YouTube doesn't provide us with the signature.
        js_url = "http:" + video_data.get("assets", {}).get("js")

        # Just make these easily accessible as variables.
        stream_map = video_data.get("args", {}).get("stream_map")
        video_urls = stream_map.get("url")

        # For each video url, identify the quality profile and add it to list
        # of available videos.
        for idx, url in enumerate(video_urls):
            log.debug("attempting to get quality profile from url: %s", url)
            try:
                itag, quality_profile = self._get_quality_profile_from_url(url)
                if not quality_profile:
                    log.warn("unable to identify profile for itag=%s", itag)
                    continue
            except (TypeError, KeyError) as e:
                log.exception("passing on exception %s", e)
                continue

            # Check if we have the signature, otherwise we'll need to get the
            # cipher from the js.
            if "signature=" not in url:
                log.debug('signature not in url, attempting to resolve the '
                          'cipher...')
                signature = self._get_cipher(stream_map["s"][idx], js_url)
                url = "{}&signature={}".format(url, signature)
            self._add_video(url, self.filename, **quality_profile)

    def get(self, extension=None, resolution=None, profile=None):
        """Gets a single video given a file extention (and/or resolution
        and/or quality profile).

        :param str extention:
            The desired file extention (e.g.: mp4, flv).
        :param str resolution:
            The desired video broadcasting standard (e.g.: 720p, 1080p)
        :param str profile:
            The desired quality profile (this is subjective, I don't recommend
            using it).
        """
        result = []
        for v in self.get_videos():
            if extension and v.extension != extension:
                continue
            elif resolution and v.resolution != resolution:
                continue
            elif profile and v.profile != profile:
                continue
            else:
                result.append(v)
        matches = len(result)
        if matches <= 0:
            return DoesNotExist("No videos met this criteria.")
        elif matches == 1:
            return result[0]
        else:
            raise MultipleObjectsReturned("Multiple videos met this criteria.")

    def filter(self, extension=None, resolution=None, profile=None):
        """Gets a filtered list of videos given a file extention and/or
        resolution and/or quality profile.

        :param str extention:
            The desired file extention (e.g.: mp4, flv).
        :param str resolution:
            The desired video broadcasting standard (e.g.: 720p, 1080p)
        :param str profile:
            The desired quality profile (this is subjective, I don't recommend
            using it).
        """
        results = []
        for v in self.get_videos():
            if extension and v.extension != extension:
                continue
            elif resolution and v.resolution != resolution:
                continue
            elif profile and v.profile != profile:
                continue
            else:
                results.append(v)
        return results

    def get_video_data(self):
        """Gets the page and extracts out the video data."""
        # Reset the filename incase it was previously set.
        self.title = None
        response = urlopen(self.url)
        if not response:
            raise PytubeError("Unable to open url: {}".format(self.url))

        html = response.read().decode("utf-8")
        if "og:restrictions:age" in html:
            raise AgeRestricted("Age restricted video. Unable to download "
                                "without being signed in.")

        # Extract out the json data from the html response body.
        json_object = self._get_json_data(html)

        # Here we decode the stream map and bundle it into the json object. We
        # do this just so we just can return one object for the video data.
        encoded_stream_map = json_object.get("args", {}).get(
            "url_encoded_fmt_stream_map")
        json_object['args']['stream_map'] = self._parse_stream_map(
            encoded_stream_map)
        return json_object

    def _parse_stream_map(self, blob):
        """A modified version of ``urlparse.parse_qs`` that's able to decode
        YouTube's stream map.

        :param str blob:
            An encoded blob of text containing the stream map data.
        """
        dct = {
            "itag": [],
            "url": [],
            "quality": [],
            "fallback_host": [],
            "s": [],
            "type": []
        }

        # Split the comma separated videos.
        videos = blob.split(",")

        # Unquote the characters and split to parameters.
        videos = [video.split("&") for video in videos]

        # Split at the equals sign so we can break this key value pairs and
        # toss it into a dictionary.
        for video in videos:
            for kv in video:
                key, value = kv.split("=")
                dct.get(key, []).append(unquote(value))
        log.debug('decoded stream map: %s', dct)
        return dct

    def _get_json_data(self, html):
        """Extract the json out from the html.

        :param str html:
            The raw html of the page.
        """
        # 18 represents the length of "ytplayer.config = ".
        start = html.find("ytplayer.config = ") + 18
        html = html[start:]

        offset = self._get_json_offset(html)
        if not offset:
            raise PytubeError("Unable to extract json.")
        return json.loads(html[:offset])

    def _get_json_offset(self, html):
        """Find where the json object starts.

        :param str html:
            The raw html of the YouTube page.
        """
        brackets = []
        index = 1
        # Determine the offset by pushing/popping brackets until all
        # js expressions are closed.
        for idx, ch in enumerate(html):
            if ch == "{":
                brackets.append("}")
            elif ch == "}":
                brackets.pop()
                if len(brackets) == 0:
                    break
        else:
            raise PytubeError("Unable to determine json offset.")
        return index + idx

    def _get_cipher(self, signature, url):
        """Gets the signature using the cipher.

        :param str signature:
            The url signature.
        :param str url:
            The url of the javascript file.
        """
        reg_exp = re.compile(r'\.sig\|\|([a-zA-Z0-9$]+)\(')
        if not self._js_code:
            response = urlopen(url)
            if not response:
                raise PytubeError("Unable to open url: {}".format(self.url))
            self._js_code = response.read().decode("utf-8")
        try:
            matches = reg_exp.search(self._js_code)
            if matches:
                # Return the first matching group.
                func = next(g for g in matches.groups() if g is not None)
            # Load js into JS Python interpreter.
            jsi = JSInterpreter(self._js_code)
            initial_function = jsi.extract_function(func)
            return initial_function([signature])
        except Exception as e:
            raise CipherError("Couldn't cipher the signature. Maybe YouTube "
                              "has changed the cipher algorithm. Notify this "
                              "issue on GitHub: {}".format(e))
        return False

    def _get_quality_profile_from_url(self, video_url):
        """Gets the quality profile given a video url. Normally we would just
        use ``urlparse`` since itags are represented as a get parameter, but
        YouTube doesn't pass a properly encoded url.

        :param str video_url:
            The malformed url-encoded video_url.
        """
        reg_exp = re.compile('itag=(\d+)')
        itag = reg_exp.findall(video_url)
        if itag and len(itag) == 1:
            itag = int(itag[0])
            # Given an itag, refer to the YouTube quality profiles to get the
            # properties (media type, resolution, etc.) of the video.
            quality_profile = YT_QUALITY_PROFILES.get(itag)
            if not quality_profile:
                return itag, None
            # Here we just combine the quality profile keys to the
            # corresponding quality profile, referenced by the itag.
            return itag, dict(zip(YT_QUALITY_PROFILE_KEYS, quality_profile))
        if not itag:
            raise PytubeError("Unable to get encoding profile, no itag found.")
        elif len(itag) > 1:
            log.warn("Multiple itags found: %s", itag)
            raise PytubeError("Unable to get encoding profile, multiple itags "
                              "found.")
        return False

    def _add_video(self, url, filename, **kwargs):
        """Adds new video object to videos.

        :param str url:
            The signed url to the video.
        :param str filename:
            The filename for the video.
        :param kwargs:
            Additional properties to set for the video object.
        """
        video = Video(url, filename, **kwargs)
        self._videos.append(video)
        self._videos.sort()
        return True
<|endoftext|>