<commit_before>
# Copyright 2013 Virantha Ekanayake All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
import shutil
import hashlib
import time
import sys

from pypdfocr_filer import PyFiler

import functools

from evernote.api.client import EvernoteClient
import evernote.edam.type.ttypes as Types
import evernote.edam.userstore.constants as UserStoreConstants
from evernote.edam.error.ttypes import EDAMUserException
from evernote.edam.error.ttypes import EDAMSystemException
from evernote.edam.error.ttypes import EDAMNotFoundException
from evernote.edam.error.ttypes import EDAMErrorCode

"""
    Implementation of a filer class
        -> Files documents to Evernote notebooks (each document becomes a new note)
"""
class en_handle(object):
    """ Generic exception handler for Evernote actions
    """
    def __init__(self, f):
        # f is the method being decorated, so save it so we can call it later!
        self.f = f
        functools.update_wrapper(self, f)

    def __get__(self, instance, owner):
        # Save a ptr to the object being decorated
        self.cls = owner
        self.obj = instance
        return self.__call__

    def __call__(self, *args, **kwargs):
        # The actual meat of the decorator

        # Call the original method being decorated
        retryCount = 3
        retry_auth = False
        msg = "EVERNOTE ERROR: %s"
        r = None
        while retryCount > 0:
            try:
                retryCount -= 1
                if retry_auth:
                    logging.debug("Retrying")
                    self.obj._connect_to_evernote(self.obj.dictUserInfo)
                retry_auth = False
                logging.debug("executing user function")
                r = self.f.__call__(self.obj, *args, **kwargs)
                break
            except EDAMUserException as e:
                err = e.errorCode
                c = EDAMErrorCode
                if err == c.AUTH_EXPIRED or err == c.DATA_REQUIRED:
                    logging.debug(msg % "Authorization expired, retrying...")
                    retry_auth = True
                    time.sleep(3)
                else:
                    logging.debug(msg % ("Unhandled error %s:%s" % (c._VALUES_TO_NAMES[err], e.parameter)))
        return r

class PyFilerEvernote(PyFiler):

    def get_target_folder(self):
        return self._target_folder
    def set_target_folder (self, target_folder):
        """ Override this to make sure we only have the basename"""
        print("Setting target_folder %s" % target_folder)
        if target_folder:
            self._target_folder = os.path.basename(target_folder)
        else:
            self._target_folder = target_folder

    target_folder = property(get_target_folder, set_target_folder)

    def get_default_folder (self):
        """ Override this to make sure we only have the basename"""
        return self._default_folder

    def set_default_folder (self, default_folder):
        """ Override this to make sure we only have the basename"""
        if default_folder:
            self._default_folder = os.path.basename(default_folder)
        else:
            self._default_folder = default_folder

    default_folder = property(get_default_folder, set_default_folder)

    def __init__(self, dev_token):
        self.target_folder = None
        self.default_folder = None
        self.original_move_folder = None
        self.folder_targets = {}
        self.dictUserInfo = { 'dev_token': dev_token }
        self._connect_to_evernote(self.dictUserInfo)

    def _connect_to_evernote(self, dictUserInfo):
        """
            Establish a connection to evernote and authenticate.

            :param dictUserInfo: Dict of user info like user/passwrod.  For now, just the dev token
            :returns success: Return wheter connection succeeded
            :rtype bool:
        """
        print("Authenticating to Evernote")
        dev_token = dictUserInfo['dev_token']
        logging.debug("Authenticating using token %s" % dev_token)
        user = None
        try:
            self.client = EvernoteClient(token=dev_token, sandbox=False)
            self.user_store = self.client.get_user_store()
            user = self.user_store.getUser()
        except EDAMUserException as e:
            err = e.errorCode
            print("Error attempting to authenticate to Evernote: %s - %s" % (EDAMErrorCode._VALUES_TO_NAMES[err], e.parameter))
        except EDAMSystemException as e:
            err = e.errorCode
            print("Error attempting to authenticate to Evernote: %s - %s" % (EDAMErrorCode._VALUES_TO_NAMES[err], e.message))
            sys.exit(-1)

        if user:
            print("Authenticated to evernote as user %s" % user.username)
        return True

    def add_folder_target(self, folder, keywords):
        assert folder not in self.folder_targets, "Target folder already defined! (%s)" % (folder)
        self.folder_targets[folder] = keywords

    def file_original(self, original_filename):
        """
            Just file it to the local file system (don't upload to evernote)
        """
        if not self.original_move_folder:
            logging.debug("Leaving original untouched")
            return original_filename

        tgt_path = self.original_move_folder
        logging.debug("Moving original %s to %s" % (original_filename, tgt_path))
        tgtfilename = os.path.join(tgt_path, os.path.basename(original_filename))
        tgtfilename = self._get_unique_filename_by_appending_version_integer(tgtfilename)

        shutil.move(original_filename, tgtfilename)
        return tgtfilename

    @en_handle
    def _get_notebooks(self):
        note_store = self.client.get_note_store()
        notebooks = note_store.listNotebooks()
        return {n.name:n for n in notebooks}

    @en_handle
    def _create_notebook(self, notebook):
        note_store = self.client.get_note_store()
        return note_store.createNotebook(notebook)

    def _update_notebook(self, notebook):
        note_store = self.client.get_note_store()
        note_store.updateNotebook(notebook)
        return

    @en_handle
    def _check_and_make_notebook(self, notebook_name):
        """
            Weird.
            :returns notebook: New or existing notebook object
            :rtype Types.Notebook:
        """
        # Get the noteStore
        #note_store = self.client.get_note_store()
        #notebooks = note_store.listNotebooks()
        #notebooks = {n.name:n for n in notebooks}
        notebooks = self._get_notebooks()
        if notebook_name in notebooks:
            notebook = notebooks[notebook_name]
            if notebook.stack != self.target_folder:
                notebook.stack = self.target_folder
                self._update_notebook(notebook)
            return notebook
        else:
            # Need to create a new notebook
            notebook = Types.Notebook()
            notebook.name = notebook_name
            notebook.stack = self.target_folder
            notebook = self._create_notebook(notebook)
            #notebook = note_store.createNotebook(notebook)
            return notebook

    @en_handle
    def _create_evernote_note(self, notebook, filename):
        # Create the new note
        note = Types.Note()
        note.title = os.path.basename(filename)
        note.notebookGuid = notebook.guid
        note.content = '<?xml version="1.0" encoding="UTF-8"?><!DOCTYPE en-note SYSTEM "http://xml.evernote.com/pub/enml2.dtd">'
        note.content += '<en-note>Uploaded by PyPDFOCR
'

        logging.debug("Loading PDF")
        md5 = hashlib.md5()
        with open(filename,'rb') as f:
            pdf_bytes = f.read()

        logging.debug("Calculating md5 checksum of pdf")
        md5.update(pdf_bytes)
        md5hash = md5.hexdigest()

        logging.debug("Uploading note")

        # Create the Data type for evernote that goes into a resource
        pdf_data = Types.Data()
        pdf_data.bodyHash = md5hash
        pdf_data.size = len(pdf_bytes)
        pdf_data.body = pdf_bytes

        # Add a link in the evernote boy for this content
        link = '<en-media type="application/pdf" hash="%s"/>' % md5hash
        logging.debug(link)
        note.content += link
        note.content += '</en-note>'

        resource_list = []
        pdf_resource = Types.Resource()
        pdf_resource.data = pdf_data
        pdf_resource.mime = "application/pdf"
        # TODO: Enable filename
        #pdf_resource.attributes.filename = filename
        resource_list.append(pdf_resource)

        note.resources = resource_list

        return note

    def move_to_matching_folder(self, filename, foldername):
        """
            Use the evernote API to create a new note:

            #. Make the notebook if it doesn't exist (:func:`_check_and_make_notebook`)
            #. Create the note (:func:`_create_evernote_note`)
            #. Upload note using API

        """
        assert self.target_folder != None
        assert self.default_folder != None

        if not foldername:
            logging.info("[DEFAULT] %s --> %s" % (filename, self.default_folder))
            foldername = self.default_folder
        else:
            logging.info("[MATCH] %s --> %s" % (filename, foldername))

        # Check if the evernote notebook exists
        print ("Checking for notebook named %s" % foldername)
        notebook = self._check_and_make_notebook(foldername)
        print("Uploading %s to %s" % (filename, foldername))

        note = self._create_evernote_note(notebook, filename)

        # Store the note in evernote
        note_store = self.client.get_note_store()
        note = note_store.createNote(note)
        os.remove(filename)

        return "%s/%s" % (notebook.name, note.title)

if __name__ == '__main__': # pragma: no cover
    logging.basicConfig(level=logging.DEBUG, format='%(message)s')
    logging.basicConfig(level=logging.INFO, format='%(message)s')
    p = PyFilerEvernote()
    p.add_folder_target("auto", ['dmv'])
    p.target_folder = 'myuploads'
    p.default_folder = 'default'
    p.original_move_folder = None

    p.move_to_matching_folder('../dmv/dmv_ocr.pdf', 'auto')
<commit_msg>Fixed title of pdfs uploaded to evernote<commit_after>
# Copyright 2013 Virantha Ekanayake All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
import shutil
import hashlib
import time
import sys

from pypdfocr_filer import PyFiler

import functools

from evernote.api.client import EvernoteClient
import evernote.edam.type.ttypes as Types
import evernote.edam.userstore.constants as UserStoreConstants
from evernote.edam.error.ttypes import EDAMUserException
from evernote.edam.error.ttypes import EDAMSystemException
from evernote.edam.error.ttypes import EDAMNotFoundException
from evernote.edam.error.ttypes import EDAMErrorCode

"""
    Implementation of a filer class
        -> Files documents to Evernote notebooks (each document becomes a new note)
"""
class en_handle(object):
    """ Generic exception handler for Evernote actions
    """
    def __init__(self, f):
        # f is the method being decorated, so save it so we can call it later!
        self.f = f
        functools.update_wrapper(self, f)

    def __get__(self, instance, owner):
        # Save a ptr to the object being decorated
        self.cls = owner
        self.obj = instance
        return self.__call__

    def __call__(self, *args, **kwargs):
        # The actual meat of the decorator

        # Call the original method being decorated
        retryCount = 3
        retry_auth = False
        msg = "EVERNOTE ERROR: %s"
        r = None
        while retryCount > 0:
            try:
                retryCount -= 1
                if retry_auth:
                    logging.debug("Retrying")
                    self.obj._connect_to_evernote(self.obj.dictUserInfo)
                retry_auth = False
                logging.debug("executing user function")
                r = self.f.__call__(self.obj, *args, **kwargs)
                break
            except EDAMUserException as e:
                err = e.errorCode
                c = EDAMErrorCode
                if err == c.AUTH_EXPIRED or err == c.DATA_REQUIRED:
                    logging.debug(msg % "Authorization expired, retrying...")
                    retry_auth = True
                    time.sleep(3)
                else:
                    logging.debug(msg % ("Unhandled error %s:%s" % (c._VALUES_TO_NAMES[err], e.parameter)))
        return r

class PyFilerEvernote(PyFiler):

    def get_target_folder(self):
        return self._target_folder
    def set_target_folder (self, target_folder):
        """ Override this to make sure we only have the basename"""
        print("Setting target_folder %s" % target_folder)
        if target_folder:
            self._target_folder = os.path.basename(target_folder)
        else:
            self._target_folder = target_folder

    target_folder = property(get_target_folder, set_target_folder)

    def get_default_folder (self):
        """ Override this to make sure we only have the basename"""
        return self._default_folder

    def set_default_folder (self, default_folder):
        """ Override this to make sure we only have the basename"""
        if default_folder:
            self._default_folder = os.path.basename(default_folder)
        else:
            self._default_folder = default_folder

    default_folder = property(get_default_folder, set_default_folder)

    def __init__(self, dev_token):
        self.target_folder = None
        self.default_folder = None
        self.original_move_folder = None
        self.folder_targets = {}
        self.dictUserInfo = { 'dev_token': dev_token }
        self._connect_to_evernote(self.dictUserInfo)

    def _connect_to_evernote(self, dictUserInfo):
        """
            Establish a connection to evernote and authenticate.

            :param dictUserInfo: Dict of user info like user/passwrod.  For now, just the dev token
            :returns success: Return wheter connection succeeded
            :rtype bool:
        """
        print("Authenticating to Evernote")
        dev_token = dictUserInfo['dev_token']
        logging.debug("Authenticating using token %s" % dev_token)
        user = None
        try:
            self.client = EvernoteClient(token=dev_token, sandbox=False)
            self.user_store = self.client.get_user_store()
            user = self.user_store.getUser()
        except EDAMUserException as e:
            err = e.errorCode
            print("Error attempting to authenticate to Evernote: %s - %s" % (EDAMErrorCode._VALUES_TO_NAMES[err], e.parameter))
        except EDAMSystemException as e:
            err = e.errorCode
            print("Error attempting to authenticate to Evernote: %s - %s" % (EDAMErrorCode._VALUES_TO_NAMES[err], e.message))
            sys.exit(-1)

        if user:
            print("Authenticated to evernote as user %s" % user.username)
        return True

    def add_folder_target(self, folder, keywords):
        assert folder not in self.folder_targets, "Target folder already defined! (%s)" % (folder)
        self.folder_targets[folder] = keywords

    def file_original(self, original_filename):
        """
            Just file it to the local file system (don't upload to evernote)
        """
        if not self.original_move_folder:
            logging.debug("Leaving original untouched")
            return original_filename

        tgt_path = self.original_move_folder
        logging.debug("Moving original %s to %s" % (original_filename, tgt_path))
        tgtfilename = os.path.join(tgt_path, os.path.basename(original_filename))
        tgtfilename = self._get_unique_filename_by_appending_version_integer(tgtfilename)

        shutil.move(original_filename, tgtfilename)
        return tgtfilename

    @en_handle
    def _get_notebooks(self):
        note_store = self.client.get_note_store()
        notebooks = note_store.listNotebooks()
        return {n.name:n for n in notebooks}

    @en_handle
    def _create_notebook(self, notebook):
        note_store = self.client.get_note_store()
        return note_store.createNotebook(notebook)

    def _update_notebook(self, notebook):
        note_store = self.client.get_note_store()
        note_store.updateNotebook(notebook)
        return

    @en_handle
    def _check_and_make_notebook(self, notebook_name):
        """
            Weird.
            :returns notebook: New or existing notebook object
            :rtype Types.Notebook:
        """
        # Get the noteStore
        #note_store = self.client.get_note_store()
        #notebooks = note_store.listNotebooks()
        #notebooks = {n.name:n for n in notebooks}
        notebooks = self._get_notebooks()
        if notebook_name in notebooks:
            notebook = notebooks[notebook_name]
            if notebook.stack != self.target_folder:
                notebook.stack = self.target_folder
                self._update_notebook(notebook)
            return notebook
        else:
            # Need to create a new notebook
            notebook = Types.Notebook()
            notebook.name = notebook_name
            notebook.stack = self.target_folder
            notebook = self._create_notebook(notebook)
            #notebook = note_store.createNotebook(notebook)
            return notebook

    @en_handle
    def _create_evernote_note(self, notebook, filename):
        # Create the new note
        note = Types.Note()
        note.title = os.path.basename(filename)
        note.notebookGuid = notebook.guid
        note.content = '<?xml version="1.0" encoding="UTF-8"?><!DOCTYPE en-note SYSTEM "http://xml.evernote.com/pub/enml2.dtd">'
        note.content += '<en-note>Uploaded by PyPDFOCR
'

        logging.debug("Loading PDF")
        md5 = hashlib.md5()
        with open(filename,'rb') as f:
            pdf_bytes = f.read()

        logging.debug("Calculating md5 checksum of pdf")
        md5.update(pdf_bytes)
        md5hash = md5.hexdigest()

        logging.debug("Uploading note")

        # Create the Data type for evernote that goes into a resource
        pdf_data = Types.Data()
        pdf_data.bodyHash = md5hash
        pdf_data.size = len(pdf_bytes)
        pdf_data.body = pdf_bytes

        # Add a link in the evernote boy for this content
        link = '<en-media type="application/pdf" hash="%s"/>' % md5hash
        logging.debug(link)
        note.content += link
        note.content += '</en-note>'

        resource_list = []
        pdf_resource = Types.Resource()
        pdf_resource.data = pdf_data
        pdf_resource.mime = "application/pdf"
        # TODO: Enable filename
        # Make a attributes for this resource
        pdf_resource.attributes = Types.ResourceAttributes()
        pdf_resource.attributes.fileName = filename
        resource_list.append(pdf_resource)

        note.resources = resource_list

        return note

    def move_to_matching_folder(self, filename, foldername):
        """
            Use the evernote API to create a new note:

            #. Make the notebook if it doesn't exist (:func:`_check_and_make_notebook`)
            #. Create the note (:func:`_create_evernote_note`)
            #. Upload note using API

        """
        assert self.target_folder != None
        assert self.default_folder != None

        if not foldername:
            logging.info("[DEFAULT] %s --> %s" % (filename, self.default_folder))
            foldername = self.default_folder
        else:
            logging.info("[MATCH] %s --> %s" % (filename, foldername))

        # Check if the evernote notebook exists
        print ("Checking for notebook named %s" % foldername)
        notebook = self._check_and_make_notebook(foldername)
        print("Uploading %s to %s" % (filename, foldername))

        note = self._create_evernote_note(notebook, filename)

        # Store the note in evernote
        note_store = self.client.get_note_store()
        note = note_store.createNote(note)
        os.remove(filename)

        return "%s/%s" % (notebook.name, note.title)

if __name__ == '__main__': # pragma: no cover
    logging.basicConfig(level=logging.DEBUG, format='%(message)s')
    logging.basicConfig(level=logging.INFO, format='%(message)s')
    p = PyFilerEvernote()
    p.add_folder_target("auto", ['dmv'])
    p.target_folder = 'myuploads'
    p.default_folder = 'default'
    p.original_move_folder = None

    p.move_to_matching_folder('../dmv/dmv_ocr.pdf', 'auto')
<|endoftext|>


<commit_before>import pysal.core.Tables as Tables
import csv

__author__ = "Charles R Schmidt <schmidtc@gmail.com>"
__all__ = ['csvWrapper']

class csvWrapper(Tables.DataTable):

    __doc__ = Tables.DataTable.__doc__

    FORMATS = ['csv']
    MODES = ['r']

    def __init__(self, *args, **kwargs):
        """

        Examples
        --------
        >>> import pysal
        >>> file_name = pysal.examples.get_path('stl_hom.csv')
        >>> f = pysal.open(file_name,'r')
        >>> y = f.read()
        >>> f.header
        ['WKT', 'NAME', 'STATE_NAME', 'STATE_FIPS', 'CNTY_FIPS', 'FIPS', 'FIPSNO', 'HR7984', 'HR8488', 'HR8893', 'HC7984', 'HC8488', 'HC8893', 'PO7984', 'PO8488', 'PO8893', 'PE77', 'PE82', 'PE87', 'RDAC80', 'RDAC85', 'RDAC90']
        >>> f._spec
        [<type 'str'>, <type 'str'>, <type 'str'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>]

        """
        Tables.DataTable.__init__(self, *args, **kwargs)
        self.__idx = {}
        self.__len = None
        self._open()

    def __len__(self):
        return self.__len

    def _open(self):
        self.fileObj = open(self.dataPath, self.mode)
        if self.mode == 'r':
            self.dataObj = csv.reader(self.fileObj)
            data = list(self.dataObj)
            if self._determineHeader(data):
                self.header = data.pop(0)
            else:
                self.header = ['field_%d' % i for i in range(len(data[0]))]
            self._spec = self._determineSpec(data)
            self.data = data
            self.fileObj.close()
            self.__len = len(data)

    def _determineHeader(self, data):
        #head = [val.strip().replace('-','').replace('.','').isdigit() for val in data[0]]
        #if True in head: #no numbers in header!
        #    HEADER = False
        #    return HEADER
        headSpec = self._determineSpec([data[0]])
        restSpec = self._determineSpec(data[1:])
        if headSpec == restSpec:
            HEADER = False
            return HEADER
        return True

    @staticmethod
    def _determineSpec(data):
        cols = len(data[0])
        spec = []
        for j in range(cols):
            isInt = True
            isFloat = True
            for row in data:
                val = row[j]
                if not val.strip().replace('-', '').replace('.', '').isdigit():
                    isInt = False
                    isFloat = False
                    break
                else:
                    if isInt and '.' in val:
                        isInt = False
            if isInt:
                spec.append(int)
            elif isFloat:
                spec.append(float)
            else:
                spec.append(str)
        return spec

    def _read(self):
        if self.pos < len(self):
            row = self.data[self.pos]
            self.pos += 1
            return row
        else:
            return None

<commit_msg>adding Universal newline support to csvReader, fixes #235<commit_after>import pysal.core.Tables as Tables
import csv

__author__ = "Charles R Schmidt <schmidtc@gmail.com>"
__all__ = ['csvWrapper']

class csvWrapper(Tables.DataTable):

    __doc__ = Tables.DataTable.__doc__

    FORMATS = ['csv']
    READ_MODES = ['r','Ur','rU','U']
    MODES = READ_MODES[:]

    def __init__(self, *args, **kwargs):
        """

        Examples
        --------
        >>> import pysal
        >>> file_name = pysal.examples.get_path('stl_hom.csv')
        >>> f = pysal.open(file_name,'r')
        >>> y = f.read()
        >>> f.header
        ['WKT', 'NAME', 'STATE_NAME', 'STATE_FIPS', 'CNTY_FIPS', 'FIPS', 'FIPSNO', 'HR7984', 'HR8488', 'HR8893', 'HC7984', 'HC8488', 'HC8893', 'PO7984', 'PO8488', 'PO8893', 'PE77', 'PE82', 'PE87', 'RDAC80', 'RDAC85', 'RDAC90']
        >>> f._spec
        [<type 'str'>, <type 'str'>, <type 'str'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'int'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>, <type 'float'>]

        """
        Tables.DataTable.__init__(self, *args, **kwargs)
        self.__idx = {}
        self.__len = None
        self._open()

    def __len__(self):
        return self.__len

    def _open(self):
        self.fileObj = open(self.dataPath, self.mode)
        if self.mode in self.READ_MODES:
            self.dataObj = csv.reader(self.fileObj)
            data = list(self.dataObj)
            if self._determineHeader(data):
                self.header = data.pop(0)
            else:
                self.header = ['field_%d' % i for i in range(len(data[0]))]
            self._spec = self._determineSpec(data)
            self.data = data
            self.fileObj.close()
            self.__len = len(data)

    def _determineHeader(self, data):
        #head = [val.strip().replace('-','').replace('.','').isdigit() for val in data[0]]
        #if True in head: #no numbers in header!
        #    HEADER = False
        #    return HEADER
        headSpec = self._determineSpec([data[0]])
        restSpec = self._determineSpec(data[1:])
        if headSpec == restSpec:
            HEADER = False
            return HEADER
        return True

    @staticmethod
    def _determineSpec(data):
        cols = len(data[0])
        spec = []
        for j in range(cols):
            isInt = True
            isFloat = True
            for row in data:
                val = row[j]
                if not val.strip().replace('-', '').replace('.', '').isdigit():
                    isInt = False
                    isFloat = False
                    break
                else:
                    if isInt and '.' in val:
                        isInt = False
            if isInt:
                spec.append(int)
            elif isFloat:
                spec.append(float)
            else:
                spec.append(str)
        return spec

    def _read(self):
        if self.pos < len(self):
            row = self.data[self.pos]
            self.pos += 1
            return row
        else:
            return None

<|endoftext|>


<commit_before>#!/usr/bin/python2.7
# encoding: utf-8

from __future__ import division
import numpy as np
import h5py

class _load_adcp:
    """
'Variables' subset in ADCP class contains the following numpy arrays:
-----------------------------------------------------------

                  _bins = depth of measurement bins, 1D array, shape=(bins)
                 |_dir_vel = velocity direction time serie, 2D array, shape=(time,bins)
                 |_east_vel = East velocity time serie, 2D array, shape=(time,bins)
                 |_lat = latitude, float, decimal degrees
                 |_lon = lontitude, float, decimal degrees
                 |_mag_signed_vel = signed velocity time serie, 2D array, shape=(time,bins)
                 |_matlabTime = matlab time, 1D array, shape=(time)
 ADCP.Variables._|_north_vel = East velocity time serie, 2D array, shape=(time,bins)
                 |_percent_of_depth = percent of the water column measured by ADCP, float
                 |_surf = pressure at surface timeserie, 1D array, shape=(time)
                 |_ua = depth averaged velocity component timeserie, 1D array, shape=(time)
                 |_va = depth averaged velocity component timeserie, 1D array, shape=(time)
                 |_ucross = ???, 1D array, shape=(time)
                 |_ualong = ???, 1D array, shape=(time)
    """
    def __init__(self,cls, debug=False):
        if debug:
            print 'Loading variables...'
        self.percent_of_depth=0.95
        #TR: fudge factor, squeeze out the 5 top % of the water column

        #Test if load with h5py or scipy
        if not type(cls.Data)==h5py._hl.files.File:
            self.lat = cls.Data['lat']
            self.lon = cls.Data['lon']
            self.bins = cls.Data['data'].bins[:].flatten()
            self.north_vel = cls.Data['data'].north_vel[:]
            self.east_vel = cls.Data['data'].east_vel[:]
            self.vert_vel = cls.Data['data'].vert_vel[:]
            self.dir_vel = cls.Data['data'].dir_vel[:]
            self.mag_signed_vel = cls.Data['data'].mag_signed_vel[:]
            self.pressure = cls.Data['pres']
            self.surf = self.pressure.surf[:].flatten()
            self.time = cls.Data['time']
            self.matlabTime = self.time.mtime[:].flatten()
            try:
                self.ucross = cls.Data['data'].ucross[:]
                self.ualong = cls.Data['data'].ualong[:]
            except AttributeError:
                pass

        else:
            self.lat = cls.Data['lat'][0][0]
            self.lon = cls.Data['lon'][0][0]
            self.bins = cls.Data['data']['bins'][:].flatten()
            self.north_vel = cls.Data['data']['north_vel'][:].T
            self.east_vel = cls.Data['data']['east_vel'][:].T
            self.vert_vel = cls.Data['data']['vert_vel'][:].T
            self.dir_vel = cls.Data['data']['dir_vel'][:].T
            self.mag_signed_vel = cls.Data['data']['mag_signed_vel'][:].T
            self.pressure = cls.Data['pres']
            self.surf = self.pressure['surf'][:].flatten()
            self.time = cls.Data['time']
            self.matlabTime = self.time['mtime'][:].flatten()
            try:
                self.ucross = cls.Data['data']['Ucross'][:].T
                self.ualong = cls.Data['data']['Ualong'][:].T
            except KeyError:
                pass

        #Find the depth average of a variable based on percent_of_depth
        #choosen by the user. Currently only working for east_vel (u) and
        #north_vel (v)
        #TR: alaternative with percent of the depth
        ind = np.argwhere(self.bins < self.percent_of_depth * self.surf[:,None])
        #ind = np.argwhere(self.bins < self.surf[:,None])
        index = ind[np.r_[ind[1:,0] != ind[:-1,0], True]]
        try:
            data_ma_u = np.ma.array(self.east_vel,
                        mask=np.arange(self.east_vel.shape[1]) > index[:, 1, None])
            data_ma_u=np.ma.masked_array(data_ma_u,np.isnan(data_ma_u))
        except MaskError:
            data_ma_u=np.ma.masked_array(self.east_vel,np.isnan(self.east_vel))

        try:
            data_ma_v = np.ma.array(self.north_vel,
                        mask=np.arange(self.north_vel.shape[1]) > index[:, 1, None])
            data_ma_v=np.ma.masked_array(data_ma_v,np.isnan(data_ma_v))
        except MaskError:
            data_ma_v=np.ma.masked_array(self.north_vel,np.isnan(self.north_vel))

        self.ua = np.array(data_ma_u.mean(axis=1))
        self.va = np.array(data_ma_v.mean(axis=1))

        if debug:
            print '...Passed'
<commit_msg>Fixed mask in depth av adcp velocities<commit_after>#!/usr/bin/python2.7
# encoding: utf-8

from __future__ import division
import numpy as np
import h5py

class _load_adcp:
    """
'Variables' subset in ADCP class contains the following numpy arrays:
-----------------------------------------------------------

                  _bins = depth of measurement bins, 1D array, shape=(bins)
                 |_dir_vel = velocity direction time serie, 2D array, shape=(time,bins)
                 |_east_vel = East velocity time serie, 2D array, shape=(time,bins)
                 |_lat = latitude, float, decimal degrees
                 |_lon = lontitude, float, decimal degrees
                 |_mag_signed_vel = signed velocity time serie, 2D array, shape=(time,bins)
                 |_matlabTime = matlab time, 1D array, shape=(time)
 ADCP.Variables._|_north_vel = East velocity time serie, 2D array, shape=(time,bins)
                 |_percent_of_depth = percent of the water column measured by ADCP, float
                 |_surf = pressure at surface timeserie, 1D array, shape=(time)
                 |_ua = depth averaged velocity component timeserie, 1D array, shape=(time)
                 |_va = depth averaged velocity component timeserie, 1D array, shape=(time)
                 |_ucross = ???, 1D array, shape=(time)
                 |_ualong = ???, 1D array, shape=(time)
    """
    def __init__(self,cls, debug=False):
        if debug:
            print 'Loading variables...'
        self.percent_of_depth=0.95
        #TR: fudge factor, squeeze out the 5 top % of the water column

        #Test if load with h5py or scipy
        if not type(cls.Data)==h5py._hl.files.File:
            self.lat = cls.Data['lat']
            self.lon = cls.Data['lon']
            self.bins = cls.Data['data'].bins[:].flatten()
            self.north_vel = cls.Data['data'].north_vel[:]
            self.east_vel = cls.Data['data'].east_vel[:]
            self.vert_vel = cls.Data['data'].vert_vel[:]
            self.dir_vel = cls.Data['data'].dir_vel[:]
            self.mag_signed_vel = cls.Data['data'].mag_signed_vel[:]
            self.pressure = cls.Data['pres']
            self.surf = self.pressure.surf[:].flatten()
            self.time = cls.Data['time']
            self.matlabTime = self.time.mtime[:].flatten()
            try:
                self.ucross = cls.Data['data'].ucross[:]
                self.ualong = cls.Data['data'].ualong[:]
            except AttributeError:
                pass

        else:
            self.lat = cls.Data['lat'][0][0]
            self.lon = cls.Data['lon'][0][0]
            self.bins = cls.Data['data']['bins'][:].flatten()
            self.north_vel = cls.Data['data']['north_vel'][:].T
            self.east_vel = cls.Data['data']['east_vel'][:].T
            self.vert_vel = cls.Data['data']['vert_vel'][:].T
            self.dir_vel = cls.Data['data']['dir_vel'][:].T
            self.mag_signed_vel = cls.Data['data']['mag_signed_vel'][:].T
            self.pressure = cls.Data['pres']
            self.surf = self.pressure['surf'][:].flatten()
            self.time = cls.Data['time']
            self.matlabTime = self.time['mtime'][:].flatten()
            try:
                self.ucross = cls.Data['data']['Ucross'][:].T
                self.ualong = cls.Data['data']['Ualong'][:].T
            except KeyError:
                pass

        #Find the depth average of a variable based on percent_of_depth
        #choosen by the user. Currently only working for east_vel (u) and
        #north_vel (v)
        #TR: alaternative with percent of the depth
        ind = np.argwhere(self.bins < self.percent_of_depth * self.surf[:,None])
        #ind = np.argwhere(self.bins < self.surf[:,None])
        index = ind[np.r_[ind[1:,0] != ind[:-1,0], True]]
        try:
            data_ma_u = np.ma.array(self.east_vel,
                        mask=np.arange(self.east_vel.shape[1]) > index[:, 1, None])
            data_ma_u=np.ma.masked_array(data_ma_u,np.isnan(data_ma_u))
        except:  #MaskError:
            data_ma_u=np.ma.masked_array(self.east_vel,np.isnan(self.east_vel))

        try:
            data_ma_v = np.ma.array(self.north_vel,
                        mask=np.arange(self.north_vel.shape[1]) > index[:, 1, None])
            data_ma_v=np.ma.masked_array(data_ma_v,np.isnan(data_ma_v))
        except:  #MaskError:
            data_ma_v=np.ma.masked_array(self.north_vel,np.isnan(self.north_vel))

        self.ua = np.array(data_ma_u.mean(axis=1))
        self.va = np.array(data_ma_v.mean(axis=1))

        if debug:
            print '...Passed'
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# vim: set et sw=4 ts=4 sts=4 ff=unix fenc=utf8:
# Author: Binux<i@binux.me>
#         http://binux.me
# Created on 2012-12-17 11:07:19

import time
import json
import Queue
import logging
import threading
import cookie_utils
import tornado.ioloop
import tornado.httputil
import tornado.httpclient
from tornado.curl_httpclient import CurlAsyncHTTPClient
from tornado.simple_httpclient import SimpleAsyncHTTPClient
from pyspider.libs import dataurl, counter
logger = logging.getLogger('fetcher')

class MyCurlAsyncHTTPClient(CurlAsyncHTTPClient):
    def free_size(self):
        return len(self._free_list)
    def size(self):
        return len(self._curls) - self.free_size()

class MySimpleAsyncHTTPClient(SimpleAsyncHTTPClient):
    def free_size(self):
        return self.max_clients - self.size()
    def size(self):
        return len(self.active)

fetcher_output = {
        "status_code": int,
        "orig_url": str,
        "url": str,
        "headers": dict,
        "content": str,
        "cookies": dict,
        }

class Fetcher(object):
    user_agent = "pyspider/master (+http://pyspider.org/)"
    default_options = {
            'method': 'GET',
            'headers': {},
            'timeout': 120,
            }
    phantomjs_proxy = None

    def __init__(self, inqueue, outqueue, poolsize=10, proxy=None, async=True):
        self.inqueue = inqueue
        self.outqueue = outqueue

        self.poolsize = poolsize
        self._running = False
        self._quit = False
        self.proxy = proxy
        self.async = async

        if async:
            self.http_client = MyCurlAsyncHTTPClient(max_clients=self.poolsize)
        else:
            self.http_client = tornado.httpclient.HTTPClient(MyCurlAsyncHTTPClient, max_clients=self.poolsize)

        self._cnt = {
                '5m': counter.CounterManager(
                    lambda : counter.TimebaseAverageWindowCounter(30, 10)),
                '1h': counter.CounterManager(
                    lambda : counter.TimebaseAverageWindowCounter(60, 60)),
                }

    def send_result(self, type, task, result):
        """type in ('data', 'http')"""
        if self.outqueue:
            try:
                self.outqueue.put((task, result))
            except Exception, e:
                logger.exception(e)

    def fetch(self, task, callback=None):
        url = task.get('url', 'data:,')
        if callback is None:
            callback = self.send_result
        if url.startswith('data:'):
            return self.data_fetch(url, task, callback)
        elif task.get('fetch', {}).get('fetch_type') in ('js', 'phantomjs'):
            return self.phantomjs_fetch(url, task, callback)
        else:
            return self.http_fetch(url, task, callback)

    def sync_fetch(self, task):
        wait_result = threading.Condition()
        _result = {}
        def callback(type, task, result):
            wait_result.acquire()
            _result['type'] = type
            _result['task'] = task
            _result['result'] = result
            wait_result.notify()
            wait_result.release()
        self.fetch(task, callback=callback)

        wait_result.acquire()
        while 'result' not in _result:
            wait_result.wait()
        wait_result.release()
        return _result['result']

    def data_fetch(self, url, task, callback):
        self.on_fetch('data', task)
        result = {}
        result['orig_url'] = url
        result['content'] = dataurl.decode(url)
        result['headers'] = {}
        result['status_code'] = 200
        result['url'] = url
        result['cookies'] = {}
        result['time'] = 0
        result['save'] = task.get('fetch', {}).get('save')
        if len(result['content']) < 70:
            logger.info("[200] %s 0s", url)
        else:
            logger.info("[200] data:,%s...[content:%d] 0s", result['content'][:70], len(result['content']))

        callback('data', task, result)
        self.on_result('data', task, result)
        return task, result

    allowed_options = ['method', 'data', 'timeout', 'allow_redirects', 'cookies']
    def http_fetch(self, url, task, callback):
        self.on_fetch('http', task)
        fetch = dict(self.default_options)
        fetch.setdefault('url', url)
        fetch.setdefault('headers', {})
        fetch.setdefault('allow_redirects', True)
        fetch.setdefault('use_gzip', True)
        fetch['headers'].setdefault('User-Agent', self.user_agent)
        task_fetch = task.get('fetch', {})
        for each in self.allowed_options:
            if each in task_fetch:
                fetch[each] = task_fetch[each]
        fetch['headers'].update(task_fetch.get('headers', {}))

        track_headers = task.get('track', {}).get('fetch', {}).get('headers', {})
        #proxy
        if 'proxy' in task_fetch:
            if isinstance(task_fetch['proxy'], basestring):
                fetch['proxy_host'] = task_fetch['proxy'].split(":")[0]
                fetch['proxy_port'] = int(task_fetch['proxy'].split(":")[1])
            elif self.proxy and task_fetch.get('proxy', True):
                fetch['proxy_host'] = self.proxy.split(":")[0]
                fetch['proxy_port'] = int(self.proxy.split(":")[1])
        #etag
        if task_fetch.get('etag', True):
            _t = task_fetch.get('etag') if isinstance(task_fetch.get('etag'), basestring) \
                                          else track_headers.get('etag')
            if _t:
                fetch['headers'].setdefault('If-None-Match', _t)
        #last modifed
        if task_fetch.get('last_modified', True):
            _t = task_fetch.get('last_modifed') \
                        if isinstance(task_fetch.get('last_modifed'), basestring) \
                        else track_headers.get('last-modified')
            if _t:
                fetch['headers'].setdefault('If-Modifed-Since', _t)

        #fix for tornado request obj
        if 'allow_redirects' in fetch:
            fetch['follow_redirects'] = fetch['allow_redirects']
            del fetch['allow_redirects']
        if 'timeout' in fetch:
            fetch['connect_timeout'] = fetch['timeout']
            fetch['request_timeout'] = fetch['timeout']
            del fetch['timeout']
        if 'data' in fetch:
            fetch['body'] = fetch['data']
            del fetch['data']
        cookie = None
        if 'cookies' in fetch:
            cookie = fetch['cookies']
            del fetch['cookies']

        def handle_response(response):
            response.headers = final_headers
            session.extract_cookies_to_jar(request, cookie_headers)
            if response.error and not isinstance(response.error, tornado.httpclient.HTTPError):
                result = {
                        'status_code': 599,
                        'error': "%r" % response.error,
                        'content': "",
                        'time': time.time() - start_time,
                        'orig_url': url,
                        'url': url,
                        }
                callback('http', task, result)
                self.on_result('http', task, result)
                return task, result
            result = {}
            result['orig_url'] = url
            result['content'] = response.body or ''
            result['headers'] = dict(response.headers)
            result['status_code'] = response.code
            result['url'] = response.effective_url or url
            result['cookies'] = session.to_dict()
            result['time'] = time.time() - start_time
            result['save'] = task_fetch.get('save')
            if 200 <= response.code < 300:
                logger.info("[%d] %s %.2fs", response.code, url, result['time'])
            else:
                logger.warning("[%d] %s %.2fs", response.code, url, result['time'])
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

        def header_callback(line):
            line = line.strip()
            if line.startswith("HTTP/"):
                final_headers.clear()
                return
            if not line:
                return
            final_headers.parse_line(line)
            cookie_headers.parse_line(line)

        start_time = time.time()
        session = cookie_utils.CookieSession()
        cookie_headers = tornado.httputil.HTTPHeaders()
        final_headers = tornado.httputil.HTTPHeaders()
        try:
            request = tornado.httpclient.HTTPRequest(header_callback=header_callback, **fetch)
            if cookie:
                session.update(cookie)
                if 'Cookie' in request.headers:
                    del request.headers['Cookie']
                request.headers['Cookie'] = session.get_cookie_header(request)
            if self.async:
                response = self.http_client.fetch(request, handle_response)
            else:
                return handle_response(self.http_client.fetch(request))
        except tornado.httpclient.HTTPError as e:
            return handle_response(e.response)
        except Exception as e:
            raise
            result = {
                    'status_code': 599,
                    'error': '%r' % e,
                    'content': "",
                    'time': time.time() - start_time,
                    'orig_url': url,
                    'url': url,
                    }
            logger.error("[599] %s, %r %.2fs", url, e, result['time'])
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

    phantomjs_adding_options = ['js_run_at', 'js_script', 'load_images']
    def phantomjs_fetch(self, url, task, callback):
        self.on_fetch('phantomjs', task)
        if not self.phantomjs_proxy:
            result = {
                    "orig_url": url,
                    "content": "phantomjs is not enabled.",
                    "headers": {},
                    "status_code": 501,
                    "url": url,
                    "cookies": {},
                    "time": 0,
                    "save": task.get('fetch', {}).get('save')
                    }
            logger.warning("[501] %s 0s", url)
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

        request_conf = {
                'follow_redirects': False
                }

        fetch = dict(self.default_options)
        fetch.setdefault('url', url)
        fetch.setdefault('headers', {})
        task_fetch = task.get('fetch', {})
        fetch.update(task_fetch)
        if 'timeout' in fetch:
            request_conf['connect_timeout'] = fetch['timeout']
            request_conf['request_timeout'] = fetch['timeout']
        fetch['headers'].setdefault('User-Agent', self.user_agent)

        start_time = time.time()
        def handle_response(response):
            if not response:
                result = {
                        'status_code': 599,
                        'error': "599 Timeout error",
                        'content': "",
                        'time': time.time() - start_time,
                        'orig_url': url,
                        'url': url,
                        }
            else:
                try:
                    result = json.loads(response.body)
                except Exception as e:
                    result = {
                            'status_code': 599,
                            'error': '%r' % e,
                            'content': '',
                            'time': time.time() - start_time,
                            'orig_url': url,
                            'url': url,
                            }
            if result.get('status_code', 200):
                logger.info("[%d] %s %.2fs", result['status_code'], url, result['time'])
            else:
                logger.exception("[%d] %s, %r %.2fs", result['status_code'],
                        url, result['content'], result['time'])
            callback('phantomjs', task, result)
            self.on_result('phantomjs', task, result)
            return task, result

        try:
            request = tornado.httpclient.HTTPRequest(
                    url="%s" % self.phantomjs_proxy, method="POST",
                    body=json.dumps(fetch), **request_conf)
            if self.async:
                response = self.http_client.fetch(request, handle_response)
            else:
                return handle_response(self.http_client.fetch(request))
        except tornado.httpclient.HTTPError as e:
            return handle_response(e.response)
        except Exception as e:
            result = {
                    'status_code': 599,
                    'error': "%r" % e,
                    'content': '',
                    'time': time.time() - start_time,
                    'orig_url': url,
                    'url': url,
                    }
            logger.error("[599] %s, %r %.2fs", url, e, result['time'])
            callback('phantomjs', task, result)
            self.on_result('phantomjs', task, result)
            return task, result

    def run(self):
        def queue_loop():
            if not self.outqueue or not self.inqueue:
                return
            while not self._quit:
                try:
                    if self.outqueue.full():
                        break
                    if self.http_client.free_size() <= 0:
                        break
                    task = self.inqueue.get_nowait()
                    self.fetch(task)
                except Queue.Empty:
                    break
                except KeyboardInterrupt:
                    break
                except Exception, e:
                    logger.exception(e)
                    break

        tornado.ioloop.PeriodicCallback(queue_loop, 100).start()
        self._running = True
        try:
            tornado.ioloop.IOLoop.instance().start()
        except KeyboardInterrupt:
            pass

        logger.info("fetcher exiting...")

    def size(self):
        return self.http_client.size()

    def quit(self):
        self._running = False
        self._quit = True
        tornado.ioloop.IOLoop.instance().stop()

    def xmlrpc_run(self, port=24444, bind='127.0.0.1', logRequests=False):
        import umsgpack
        from SimpleXMLRPCServer import SimpleXMLRPCServer
        from xmlrpclib import Binary

        server = SimpleXMLRPCServer((bind, port), allow_none=True, logRequests=logRequests)
        server.register_introspection_functions()
        server.register_multicall_functions()

        server.register_function(self.quit, '_quit')
        server.register_function(self.size)
        def sync_fetch(task):
            result = self.sync_fetch(task)
            result = Binary(umsgpack.packb(result))
            return result
        server.register_function(sync_fetch, 'fetch')
        def dump_counter(_time, _type):
            return self._cnt[_time].to_dict(_type)
        server.register_function(dump_counter, 'counter')

        server.timeout = 0.5
        while not self._quit:
            server.handle_request()
        server.server_close()

    def on_fetch(self, type, task):
        """type in ('data', 'http')"""
        pass

    def on_result(self, type, task, result):
        """type in ('data', 'http')"""
        status_code = result.get('status_code', 599)
        if status_code != 599:
            status_code = (int(status_code) / 100 * 100)
        self._cnt['5m'].event((task.get('project'), status_code), +1)
        self._cnt['1h'].event((task.get('project'), status_code), +1)

        if type == 'http' and result.get('time'):
            content_len = len(result.get('content', ''))
            self._cnt['5m'].event((task.get('project'), 'speed'), float(content_len)/result.get('time'))
            self._cnt['1h'].event((task.get('project'), 'speed'), float(content_len)/result.get('time'))
            self._cnt['5m'].event((task.get('project'), 'time'), result.get('time'))
            self._cnt['1h'].event((task.get('project'), 'time'), result.get('time'))
<commit_msg>increase default fetcher poolsize to 100<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# vim: set et sw=4 ts=4 sts=4 ff=unix fenc=utf8:
# Author: Binux<i@binux.me>
#         http://binux.me
# Created on 2012-12-17 11:07:19

import time
import json
import Queue
import logging
import threading
import cookie_utils
import tornado.ioloop
import tornado.httputil
import tornado.httpclient
from tornado.curl_httpclient import CurlAsyncHTTPClient
from tornado.simple_httpclient import SimpleAsyncHTTPClient
from pyspider.libs import dataurl, counter
logger = logging.getLogger('fetcher')

class MyCurlAsyncHTTPClient(CurlAsyncHTTPClient):
    def free_size(self):
        return len(self._free_list)
    def size(self):
        return len(self._curls) - self.free_size()

class MySimpleAsyncHTTPClient(SimpleAsyncHTTPClient):
    def free_size(self):
        return self.max_clients - self.size()
    def size(self):
        return len(self.active)

fetcher_output = {
        "status_code": int,
        "orig_url": str,
        "url": str,
        "headers": dict,
        "content": str,
        "cookies": dict,
        }

class Fetcher(object):
    user_agent = "pyspider/master (+http://pyspider.org/)"
    default_options = {
            'method': 'GET',
            'headers': {},
            'timeout': 120,
            }
    phantomjs_proxy = None

    def __init__(self, inqueue, outqueue, poolsize=100, proxy=None, async=True):
        self.inqueue = inqueue
        self.outqueue = outqueue

        self.poolsize = poolsize
        self._running = False
        self._quit = False
        self.proxy = proxy
        self.async = async

        if async:
            self.http_client = MyCurlAsyncHTTPClient(max_clients=self.poolsize)
        else:
            self.http_client = tornado.httpclient.HTTPClient(MyCurlAsyncHTTPClient, max_clients=self.poolsize)

        self._cnt = {
                '5m': counter.CounterManager(
                    lambda : counter.TimebaseAverageWindowCounter(30, 10)),
                '1h': counter.CounterManager(
                    lambda : counter.TimebaseAverageWindowCounter(60, 60)),
                }

    def send_result(self, type, task, result):
        """type in ('data', 'http')"""
        if self.outqueue:
            try:
                self.outqueue.put((task, result))
            except Exception, e:
                logger.exception(e)

    def fetch(self, task, callback=None):
        url = task.get('url', 'data:,')
        if callback is None:
            callback = self.send_result
        if url.startswith('data:'):
            return self.data_fetch(url, task, callback)
        elif task.get('fetch', {}).get('fetch_type') in ('js', 'phantomjs'):
            return self.phantomjs_fetch(url, task, callback)
        else:
            return self.http_fetch(url, task, callback)

    def sync_fetch(self, task):
        wait_result = threading.Condition()
        _result = {}
        def callback(type, task, result):
            wait_result.acquire()
            _result['type'] = type
            _result['task'] = task
            _result['result'] = result
            wait_result.notify()
            wait_result.release()
        self.fetch(task, callback=callback)

        wait_result.acquire()
        while 'result' not in _result:
            wait_result.wait()
        wait_result.release()
        return _result['result']

    def data_fetch(self, url, task, callback):
        self.on_fetch('data', task)
        result = {}
        result['orig_url'] = url
        result['content'] = dataurl.decode(url)
        result['headers'] = {}
        result['status_code'] = 200
        result['url'] = url
        result['cookies'] = {}
        result['time'] = 0
        result['save'] = task.get('fetch', {}).get('save')
        if len(result['content']) < 70:
            logger.info("[200] %s 0s", url)
        else:
            logger.info("[200] data:,%s...[content:%d] 0s", result['content'][:70], len(result['content']))

        callback('data', task, result)
        self.on_result('data', task, result)
        return task, result

    allowed_options = ['method', 'data', 'timeout', 'allow_redirects', 'cookies']
    def http_fetch(self, url, task, callback):
        self.on_fetch('http', task)
        fetch = dict(self.default_options)
        fetch.setdefault('url', url)
        fetch.setdefault('headers', {})
        fetch.setdefault('allow_redirects', True)
        fetch.setdefault('use_gzip', True)
        fetch['headers'].setdefault('User-Agent', self.user_agent)
        task_fetch = task.get('fetch', {})
        for each in self.allowed_options:
            if each in task_fetch:
                fetch[each] = task_fetch[each]
        fetch['headers'].update(task_fetch.get('headers', {}))

        track_headers = task.get('track', {}).get('fetch', {}).get('headers', {})
        #proxy
        if 'proxy' in task_fetch:
            if isinstance(task_fetch['proxy'], basestring):
                fetch['proxy_host'] = task_fetch['proxy'].split(":")[0]
                fetch['proxy_port'] = int(task_fetch['proxy'].split(":")[1])
            elif self.proxy and task_fetch.get('proxy', True):
                fetch['proxy_host'] = self.proxy.split(":")[0]
                fetch['proxy_port'] = int(self.proxy.split(":")[1])
        #etag
        if task_fetch.get('etag', True):
            _t = task_fetch.get('etag') if isinstance(task_fetch.get('etag'), basestring) \
                                          else track_headers.get('etag')
            if _t:
                fetch['headers'].setdefault('If-None-Match', _t)
        #last modifed
        if task_fetch.get('last_modified', True):
            _t = task_fetch.get('last_modifed') \
                        if isinstance(task_fetch.get('last_modifed'), basestring) \
                        else track_headers.get('last-modified')
            if _t:
                fetch['headers'].setdefault('If-Modifed-Since', _t)

        #fix for tornado request obj
        if 'allow_redirects' in fetch:
            fetch['follow_redirects'] = fetch['allow_redirects']
            del fetch['allow_redirects']
        if 'timeout' in fetch:
            fetch['connect_timeout'] = fetch['timeout']
            fetch['request_timeout'] = fetch['timeout']
            del fetch['timeout']
        if 'data' in fetch:
            fetch['body'] = fetch['data']
            del fetch['data']
        cookie = None
        if 'cookies' in fetch:
            cookie = fetch['cookies']
            del fetch['cookies']

        def handle_response(response):
            response.headers = final_headers
            session.extract_cookies_to_jar(request, cookie_headers)
            if response.error and not isinstance(response.error, tornado.httpclient.HTTPError):
                result = {
                        'status_code': 599,
                        'error': "%r" % response.error,
                        'content': "",
                        'time': time.time() - start_time,
                        'orig_url': url,
                        'url': url,
                        }
                callback('http', task, result)
                self.on_result('http', task, result)
                return task, result
            result = {}
            result['orig_url'] = url
            result['content'] = response.body or ''
            result['headers'] = dict(response.headers)
            result['status_code'] = response.code
            result['url'] = response.effective_url or url
            result['cookies'] = session.to_dict()
            result['time'] = time.time() - start_time
            result['save'] = task_fetch.get('save')
            if 200 <= response.code < 300:
                logger.info("[%d] %s %.2fs", response.code, url, result['time'])
            else:
                logger.warning("[%d] %s %.2fs", response.code, url, result['time'])
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

        def header_callback(line):
            line = line.strip()
            if line.startswith("HTTP/"):
                final_headers.clear()
                return
            if not line:
                return
            final_headers.parse_line(line)
            cookie_headers.parse_line(line)

        start_time = time.time()
        session = cookie_utils.CookieSession()
        cookie_headers = tornado.httputil.HTTPHeaders()
        final_headers = tornado.httputil.HTTPHeaders()
        try:
            request = tornado.httpclient.HTTPRequest(header_callback=header_callback, **fetch)
            if cookie:
                session.update(cookie)
                if 'Cookie' in request.headers:
                    del request.headers['Cookie']
                request.headers['Cookie'] = session.get_cookie_header(request)
            if self.async:
                response = self.http_client.fetch(request, handle_response)
            else:
                return handle_response(self.http_client.fetch(request))
        except tornado.httpclient.HTTPError as e:
            return handle_response(e.response)
        except Exception as e:
            raise
            result = {
                    'status_code': 599,
                    'error': '%r' % e,
                    'content': "",
                    'time': time.time() - start_time,
                    'orig_url': url,
                    'url': url,
                    }
            logger.error("[599] %s, %r %.2fs", url, e, result['time'])
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

    phantomjs_adding_options = ['js_run_at', 'js_script', 'load_images']
    def phantomjs_fetch(self, url, task, callback):
        self.on_fetch('phantomjs', task)
        if not self.phantomjs_proxy:
            result = {
                    "orig_url": url,
                    "content": "phantomjs is not enabled.",
                    "headers": {},
                    "status_code": 501,
                    "url": url,
                    "cookies": {},
                    "time": 0,
                    "save": task.get('fetch', {}).get('save')
                    }
            logger.warning("[501] %s 0s", url)
            callback('http', task, result)
            self.on_result('http', task, result)
            return task, result

        request_conf = {
                'follow_redirects': False
                }

        fetch = dict(self.default_options)
        fetch.setdefault('url', url)
        fetch.setdefault('headers', {})
        task_fetch = task.get('fetch', {})
        fetch.update(task_fetch)
        if 'timeout' in fetch:
            request_conf['connect_timeout'] = fetch['timeout']
            request_conf['request_timeout'] = fetch['timeout']
        fetch['headers'].setdefault('User-Agent', self.user_agent)

        start_time = time.time()
        def handle_response(response):
            if not response:
                result = {
                        'status_code': 599,
                        'error': "599 Timeout error",
                        'content': "",
                        'time': time.time() - start_time,
                        'orig_url': url,
                        'url': url,
                        }
            else:
                try:
                    result = json.loads(response.body)
                except Exception as e:
                    result = {
                            'status_code': 599,
                            'error': '%r' % e,
                            'content': '',
                            'time': time.time() - start_time,
                            'orig_url': url,
                            'url': url,
                            }
            if result.get('status_code', 200):
                logger.info("[%d] %s %.2fs", result['status_code'], url, result['time'])
            else:
                logger.exception("[%d] %s, %r %.2fs", result['status_code'],
                        url, result['content'], result['time'])
            callback('phantomjs', task, result)
            self.on_result('phantomjs', task, result)
            return task, result

        try:
            request = tornado.httpclient.HTTPRequest(
                    url="%s" % self.phantomjs_proxy, method="POST",
                    body=json.dumps(fetch), **request_conf)
            if self.async:
                response = self.http_client.fetch(request, handle_response)
            else:
                return handle_response(self.http_client.fetch(request))
        except tornado.httpclient.HTTPError as e:
            return handle_response(e.response)
        except Exception as e:
            result = {
                    'status_code': 599,
                    'error': "%r" % e,
                    'content': '',
                    'time': time.time() - start_time,
                    'orig_url': url,
                    'url': url,
                    }
            logger.error("[599] %s, %r %.2fs", url, e, result['time'])
            callback('phantomjs', task, result)
            self.on_result('phantomjs', task, result)
            return task, result

    def run(self):
        def queue_loop():
            if not self.outqueue or not self.inqueue:
                return
            while not self._quit:
                try:
                    if self.outqueue.full():
                        break
                    if self.http_client.free_size() <= 0:
                        break
                    task = self.inqueue.get_nowait()
                    self.fetch(task)
                except Queue.Empty:
                    break
                except KeyboardInterrupt:
                    break
                except Exception, e:
                    logger.exception(e)
                    break

        tornado.ioloop.PeriodicCallback(queue_loop, 100).start()
        self._running = True
        try:
            tornado.ioloop.IOLoop.instance().start()
        except KeyboardInterrupt:
            pass

        logger.info("fetcher exiting...")

    def size(self):
        return self.http_client.size()

    def quit(self):
        self._running = False
        self._quit = True
        tornado.ioloop.IOLoop.instance().stop()

    def xmlrpc_run(self, port=24444, bind='127.0.0.1', logRequests=False):
        import umsgpack
        from SimpleXMLRPCServer import SimpleXMLRPCServer
        from xmlrpclib import Binary

        server = SimpleXMLRPCServer((bind, port), allow_none=True, logRequests=logRequests)
        server.register_introspection_functions()
        server.register_multicall_functions()

        server.register_function(self.quit, '_quit')
        server.register_function(self.size)
        def sync_fetch(task):
            result = self.sync_fetch(task)
            result = Binary(umsgpack.packb(result))
            return result
        server.register_function(sync_fetch, 'fetch')
        def dump_counter(_time, _type):
            return self._cnt[_time].to_dict(_type)
        server.register_function(dump_counter, 'counter')

        server.timeout = 0.5
        while not self._quit:
            server.handle_request()
        server.server_close()

    def on_fetch(self, type, task):
        """type in ('data', 'http')"""
        pass

    def on_result(self, type, task, result):
        """type in ('data', 'http')"""
        status_code = result.get('status_code', 599)
        if status_code != 599:
            status_code = (int(status_code) / 100 * 100)
        self._cnt['5m'].event((task.get('project'), status_code), +1)
        self._cnt['1h'].event((task.get('project'), status_code), +1)

        if type == 'http' and result.get('time'):
            content_len = len(result.get('content', ''))
            self._cnt['5m'].event((task.get('project'), 'speed'), float(content_len)/result.get('time'))
            self._cnt['1h'].event((task.get('project'), 'speed'), float(content_len)/result.get('time'))
            self._cnt['5m'].event((task.get('project'), 'time'), result.get('time'))
            self._cnt['1h'].event((task.get('project'), 'time'), result.get('time'))
<|endoftext|>