<commit_before>import urllib2
import os, sys
import ConfigParser
import datetime
import socket
import json

config = ConfigParser.RawConfigParser()
config.read(os.path.split(os.path.dirname(__file__))[0]+'/conf/config.ini')

jolokia_url = config.get('Kafka', 'jolokia')
hostname = socket.getfqdn()
cluster_name = config.get('SelfConfig', 'cluster_name')
check_type = 'kafka'

def run_kafka():
    try:
        sys.path.append(os.path.split(os.path.dirname(__file__))[0]+'/lib')
        push = __import__('pushdata')
        jsondata=push.JonSon()
        jsondata.create_data()
        push = __import__('pushdata')
        value_rate= __import__('record_rate')
        rate=value_rate.ValueRate()
        timestamp = int(datetime.datetime.now().strftime("%s"))
        sys.path.append(os.path.split(os.path.dirname(__file__))[0]+'/lib')
        jolo_mbeans=('java.lang:type=Memory',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesRejectedPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=FailedFetchRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=FailedProduceRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=TotalFetchRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=TotalProduceRequestsPerSec',
                     'kafka.server:type=Fetch',
                     'kafka.server:type=Produce'
                     )
        for beans in jolo_mbeans:
            jolo_url=urllib2.urlopen(jolokia_url+'/'+beans, timeout=15).read()
            jolo_json = json.loads(jolo_url)
            jolo_keys = jolo_json['value']
            if beans == 'java.lang:type=Memory':
                metr_name=('used', 'committed')
                heap_type=('NonHeapMemoryUsage', 'HeapMemoryUsage')
                for heap in heap_type:
                    for metr in metr_name:
                        if heap == 'NonHeapMemoryUsage':
                            key='kafka_nonheap_'+ metr
                            mon_values=jolo_keys[heap][metr]
                            jsondata.gen_data(key, timestamp, mon_values, push.hostname, check_type, cluster_name)
                        else:
                            key='kafka_heap_'+ metr
                            mon_values=jolo_keys[heap][metr]
                            jsondata.gen_data(key, timestamp, mon_values, push.hostname, check_type, cluster_name)
            elif beans == 'kafka.server:type=Fetch':
                value= jolo_keys['queue-size']
                name = 'kafka_'+beans.split('=')[1]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)

            elif beans == 'kafka.server:type=Produce':
                value= jolo_keys['queue-size']
                name = 'kafka_'+beans.split('=')[1]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)
#                    jsondata.gen_data(name, timestamp, value, push.hostname, check_type, cluster_name)
            else:
                value= jolo_keys['Count']
                name = 'kafka_'+beans.split('=')[2].split(',')[0]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)
        jsondata.put_json()
        jsondata.truncate_data()

    except Exception as e:
        push = __import__('pushdata')
        push.print_error(__name__ , (e))
        pass

<commit_msg>Update Kafka check<commit_after>import urllib2
import os, sys
import ConfigParser
import datetime
import socket
import json

config = ConfigParser.RawConfigParser()
config.read(os.path.split(os.path.dirname(__file__))[0]+'/conf/config.ini')

jolokia_url = config.get('Kafka', 'jolokia')
hostname = socket.getfqdn()
cluster_name = config.get('SelfConfig', 'cluster_name')
check_type = 'kafka'

def run_kafka():
    try:
        sys.path.append(os.path.split(os.path.dirname(__file__))[0]+'/lib')
        push = __import__('pushdata')
        jsondata=push.JonSon()
        jsondata.create_data()
        push = __import__('pushdata')
        value_rate= __import__('record_rate')
        rate=value_rate.ValueRate()
        timestamp = int(datetime.datetime.now().strftime("%s"))
        sys.path.append(os.path.split(os.path.dirname(__file__))[0]+'/lib')
        jolo_mbeans=('java.lang:type=Memory',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=BytesRejectedPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=FailedFetchRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=FailedProduceRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=TotalFetchRequestsPerSec',
                     'kafka.server:type=BrokerTopicMetrics,name=TotalProduceRequestsPerSec',
                     'kafka.server:type=Fetch',
                     'kafka.server:type=Produce',
                     'kafka.network:name=RequestsPerSec,request=FetchConsumer,type=RequestMetrics',
                     'kafka.network:name=RequestsPerSec,request=Produce,type=RequestMetrics'
                     )
        for beans in jolo_mbeans:
            jolo_url=urllib2.urlopen(jolokia_url+'/'+beans, timeout=15).read()
            jolo_json = json.loads(jolo_url)
            jolo_keys = jolo_json['value']
            if beans == 'java.lang:type=Memory':
                metr_name=('used', 'committed')
                heap_type=('NonHeapMemoryUsage', 'HeapMemoryUsage')
                for heap in heap_type:
                    for metr in metr_name:
                        if heap == 'NonHeapMemoryUsage':
                            key='kafka_nonheap_'+ metr
                            mon_values=jolo_keys[heap][metr]
                            jsondata.gen_data(key, timestamp, mon_values, push.hostname, check_type, cluster_name)
                        else:
                            key='kafka_heap_'+ metr
                            mon_values=jolo_keys[heap][metr]
                            jsondata.gen_data(key, timestamp, mon_values, push.hostname, check_type, cluster_name)
            elif beans == 'kafka.server:type=Fetch':
                value= jolo_keys['queue-size']
                name = 'kafka_'+beans.split('=')[1]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)

            elif beans == 'kafka.server:type=Produce':
                value= jolo_keys['queue-size']
                name = 'kafka_'+beans.split('=')[1]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)
#                    jsondata.gen_data(name, timestamp, value, push.hostname, check_type, cluster_name)

            elif beans == 'kafka.network:name=RequestsPerSec,request=FetchConsumer,type=RequestMetrics':
                minute_value = jolo_keys['OneMinuteRate']
                value = minute_value / 60
                key = 'kafka_RequestMetrics_FetchConsumer'
                jsondata.gen_data(key, timestamp, value, push.hostname, check_type, cluster_name)
            elif beans == 'kafka.network:name=RequestsPerSec,request=Produce,type=RequestMetrics':
                minute_value = jolo_keys['OneMinuteRate']
                value = minute_value / 60
                key = 'kafka_RequestMetrics_Produce'
                jsondata.gen_data(key, timestamp, value, push.hostname, check_type, cluster_name)

            else:
                value= jolo_keys['Count']
                name = 'kafka_'+beans.split('=')[2].split(',')[0]
                values_rate = rate.record_value_rate(name, value, timestamp)
                if values_rate >= 0:
                    jsondata.gen_data(name, timestamp, values_rate, push.hostname, check_type, cluster_name)
        jsondata.put_json()
        jsondata.truncate_data()

    except Exception as e:
        push = __import__('pushdata')
        push.print_error(__name__ , (e))
        pass

<|endoftext|>


<commit_before>from time import gmtime
import re

from util import hook, http

PRE_DB_SEARCH_URL = 'https://pr3.us/search.php'
PRE_DB_RE_NOT_FOUND = re.compile('^Nothing found for: .+$')

def get_predb_release(release_name):
    timestamp = gmtime()

    try:
        # Without accept headers, the underlying
        # site falls apart for some reason
        h = http.get_html(
            PRE_DB_SEARCH_URL,
            search=release_name,
            ts=timestamp,
            pretimezone=0,
            timezone=0,
            headers={
                'Accept': '*/*',
                'Accept-Language': 'en-US,en;q=0.5'
            }
        )
    except http.HTTPError:
        return None

    results = h.xpath("//tr")

    if not results:
        return None

    first_result = results[0]

    if PRE_DB_RE_NOT_FOUND.match(first_result.text_content()):
        return None

    section_field, name_field, _, size_field, ts_field, _ = \
        first_result.xpath('./td')

    date, time = ts_field.text_content().split()

    return {
        'date': date,
        'time': time,
        'section': section_field.text_content(),
        'name': name_field.text_content().strip(),
        'size': size_field.text_content()
    }

@hook.command
def predb(inp):
    """.predb <query> -- searches scene releases"""

    release = get_predb_release(inp)

    if not release:
        return 'zero results'

    if release['size']:
        release['size'] = ' - %s' % release['size']

    return '{date} - {section} - {name}{size}'.format(**release)
<commit_msg>not every predb type has all required fields<commit_after>from time import gmtime
import re

from util import hook, http

PRE_DB_SEARCH_URL = 'https://pr3.us/search.php'
PRE_DB_RE_NOT_FOUND = re.compile('^Nothing found for: .+$')

def get_predb_release(release_name):
    timestamp = gmtime()

    try:
        # Without accept headers, the underlying
        # site falls apart for some reason
        h = http.get_html(
            PRE_DB_SEARCH_URL,
            search=release_name,
            ts=timestamp,
            pretimezone=0,
            timezone=0,
            headers={
                'Accept': '*/*',
                'Accept-Language': 'en-US,en;q=0.5'
            }
        )
    except http.HTTPError:
        return None

    results = h.xpath("//tr")

    if not results:
        return None

    first_result = results[0]

    if PRE_DB_RE_NOT_FOUND.match(first_result.text_content()):
        return None

    section_field, name_field, _, size_field, ts_field = \
        first_result.xpath('./td')[:5]

    date, time = ts_field.text_content().split()

    return {
        'date': date,
        'time': time,
        'section': section_field.text_content(),
        'name': name_field.text_content().strip(),
        'size': size_field.text_content()
    }

@hook.command
def predb(inp):
    """.predb <query> -- searches scene releases"""

    release = get_predb_release(inp)

    if not release:
        return 'zero results'

    if release['size']:
        release['size'] = ' - %s' % release['size']

    return '{date} - {section} - {name}{size}'.format(**release)
<|endoftext|>


<commit_before>import json
import re

from bs4 import BeautifulSoup

from glob import glob
from pyspark import SparkConf, SparkContext
from pyspark.sql import SparkSession
from pyspark.ml.feature import CountVectorizer
from pyspark.ml.feature import StopWordsRemover
from pyspark.ml.feature import Word2Vec
from pyspark.ml.clustering import KMeans

def readguide(guide):
    with open(guide, 'r') as myfile:
        return myfile.read()

def extract_text(content_str):
    content = json.loads(content_str)

    soup = BeautifulSoup(content[2], "lxml")

    # kill all script and style elements
    for script in soup(["script", "style"]):
        script.extract()

    headerdiv = soup.find("div", {"id": "s-lg-guide-header-info"})
    maindiv = soup.find("div", {"id": "s-lg-guide-main"})

    if headerdiv and maindiv:
        txt = " ".join(headerdiv.strings).replace('\n', ' ') + " ".join(maindiv.strings).replace('\n', ' ')
        # remove non-alpha/numeric
        regex = re.compile('[^a-zA-Z]')
        txt = regex.sub(" ", txt)
        # remove whitespace
        txt = ' '.join(txt.split())
        if isinstance(txt, str):
            #return [content[0], content[1], txt.split(" ")]
            return [content[0], content[1], txt]
        else:
            return None

    else:
        return None

conf = SparkConf().setMaster("local").setAppName("My App")
sc = SparkContext(conf=conf)
spark = SparkSession(sc)

guides = glob("data/guides/*")

# create an RDD consisting of a list of paths to .json files
#guidesRDD = sc.parallelize(guides[0:10])
guidesRDD = sc.parallelize(guides)

# create an RDD containing html documents
documentsRDD = guidesRDD.map(lambda guide: readguide(guide))

# create RDD containg plain text derived from html, guide id, and page_id
textsRDD = documentsRDD.map(lambda document: extract_text(document))
#print(" ----------- textsRDD.count()", textsRDD.count())

filteredRDD = textsRDD.filter(lambda x: x is not None)
#print(" ----------- raw.count()", raw.count())

#print(textsRDD.take(1000))

rawDFrame = filteredRDD.toDF(["guide_id", "page_id", "words"])

#textsDFrame = raw.filter("words != ' '")
#rawDFrame.show(1000)

#textsDFrame.printSchema()

# save to parquet
rawDFrame.select("guide_id", "page_id", "words").write.save("data/libguides_bow.parquet", format="parquet")
<commit_msg>Use df.withColumn()<commit_after>import json
import re

from bs4 import BeautifulSoup

from glob import glob
from pyspark import SparkConf, SparkContext
from pyspark.sql import SparkSession
from pyspark.ml.feature import CountVectorizer
from pyspark.ml.feature import StopWordsRemover
from pyspark.ml.feature import Word2Vec
from pyspark.ml.clustering import KMeans
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

def extract_text(content_str):

    if content_str is not None:
        soup = BeautifulSoup(content_str, "lxml")

        # kill all script and style elements
        for script in soup(["script", "style"]):
            script.extract()

        headerdiv = soup.find("div", {"id": "s-lg-guide-header-info"})
        maindiv = soup.find("div", {"id": "s-lg-guide-main"})

        if headerdiv and maindiv:
            txt = " ".join(headerdiv.strings).replace('\n', ' ') + " ".join(maindiv.strings).replace('\n', ' ')
            # remove non-alpha/numeric
            regex = re.compile('[^a-zA-Z]')
            txt = regex.sub(" ", txt)
            # remove whitespace
            txt = ' '.join(txt.split())
            if isinstance(txt, str):
                return txt
            else:
                return None

        else:
            return None
    else:
        return None

conf = SparkConf().setMaster("local").setAppName("My App")
sc = SparkContext(conf=conf)
spark = SparkSession(sc)

# Define udf
from pyspark.sql.functions import udf
udfextract_text=udf(extract_text, StringType())

raw = spark.read.load("data/all_guides.parquet")
raw.show(10)

print(" ---------- done load raw")

#textsDF = raw.limit(10).withColumn("words", udfextract_text(raw.content))
textsDF = raw.withColumn("words", udfextract_text(raw.content))

print(" ---------- done textDF")
textsDF.show(10)

print(" -------- textsDF.count(): ", textsDF.count())

#not_nulDF = textsDF.filter("words is not null")
#print(" -------- not_nulDF.count(): ", not_nulDF.count())

# save to parquet
textsDF.write.save("data/libguides_bow.parquet", format="parquet")
<|endoftext|>


<commit_before>"""Version info"""
short_version = '0.1'
version = '0.1.61'
<commit_msg>Updating version for development<commit_after>"""Version info"""
short_version = '0.1'
version = '0.1.7-dev'
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
# Copyright 2014-2015 Luc Saffre
# This file is part of Lino Cosi.
#
# Lino Cosi is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Lino Cosi is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with Lino Cosi.  If not, see
# <http://www.gnu.org/licenses/>.

"""
Database models for `lino_cosi.lib.sepa`.

"""

from __future__ import unicode_literals
import logging
from pprint import pformat
import glob
import os
from django.db import models
from django.core.exceptions import MultipleObjectsReturned
from lino.api import dd, _, rt
from lino.core.utils import ChangeWatcher
from .camt import CamtParser
from .fields import IBANField, BICField
from .utils import belgian_nban_to_iban_bic, iban2bic
import time

logger = logging.getLogger(__name__)

class ImportStatements(dd.Action):
    """Import the .xml files found in the directory specified at
    :attr:`import_statements_path
    <lino_cosi.lib.sepa.Plugin.import_statements_path>`.

    End-users invoke this via the menu command :menuselection:`SEPA
    --> Import SEPA`.

    When a file has been successfully imported, Lino deletes it.

    It might happen that an .xml file accidentally gets downloaded a
    second time. Lino does not create these statements again.

    """
    label = _("Import SEPA")
    http_method = 'POST'
    select_rows = False

    def get_view_permission(self, profile):
        """Make it invisible when :attr:`import_statements_path
        <lino_cosi.lib.sepa.Plugin.import_statements_path>` is empty.

        """
        if not dd.plugins.sepa.import_statements_path:
            return False
        return super(ImportStatements, self).get_view_permission(profile)

    def run_from_ui(self, ar):
        pth = dd.plugins.sepa.import_statements_path
        wc = os.path.join(pth, '*.xml')
        dd.logger.info("Importing SEPA statements from %s...", wc)
        count = 0
        if pth:
            for filename in glob.iglob(wc):
                self.import_file(ar, filename)
                count += 1
            msg = "{0} xml files would have been imported.".format(count)
            dd.logger.info(msg)
            return ar.success(msg, alert=_("Success"))
        msg = "No import_statements_path configured."
        return ar.error(msg, alert=_("Error"))

    def import_file(self, ar, filename):
        Account = rt.modules.sepa.Account

        # Parse a CAMT053 XML file.
        parser = CamtParser()
        data_file = open(filename, 'rb').read()
        num = 0
        failed_statements = {}
        try:
            dd.logger.info("Parsing %s with camt.", filename)
            res = parser.parse(data_file)
            if res is None:
                raise Exception("res is None")
            for _statement in res:
                iban = _statement['account_number']
                if iban is None:
                    # msg = "Statement without account number : {0}"
                    failed_statements[num] = "IBAN Not found"
                    continue
                    # raise Exception(msg.format(pformat(_statement)))
                try:
                    account = Account.objects.get(iban=iban)
                except Account.DoesNotExist:
                    account = Account(iban=iban)
                    account.full_clean()
                    account.save()
                except MultipleObjectsReturned:
                    msg = "Found more than one account with IBAN {0}"
                    dd.logger.warning(msg.format(iban))
                s = Statement(account=account,
                              date=_statement['date'].strftime("%Y-%m-%d"),
                              date_done=time.strftime("%Y-%m-%d"),
                              statement_number=_statement['name'],
                              balance_end=_statement['balance_end'],
                              balance_start=_statement['balance_start'],
                              balance_end_real=_statement['balance_end_real'],
                              currency_code=_statement['currency_code'])
                s.save()
                for _movement in _statement['transactions']:
                    # TODO :check if the movement is already imported.
                    if not Movement.objects.filter(
                            unique_import_id=_movement['unique_import_id']).exists():
                        _ref = _movement.get('ref', '') or ''
                        if _movement.remote_account:
                            try:
                                _bank_account = Account.objects.get(iban=_movement.remote_account)
                            except Account.DoesNotExist:
                                _bank_account = Account(iban=_movement.remote_account)
                                _bank_account.full_clean()
                                _bank_account.save()
                            m = Movement(statement=s,
                                         unique_import_id=_movement['unique_import_id'],
                                         movement_date=_movement['date'],
                                         amount=_movement['amount'],
                                         partner_name=_movement.remote_owner,
                                         ref=_ref,
                                         bank_account=_bank_account)
                            m.save()
                num += 1

        except ValueError:
            dd.logger.info("Statement file was not a camt file.")

        if len(failed_statements) == 0:
            msg = "Imported {0} statemenmts from file {1}.".format(num, filename)
            dd.logger.info(msg)
            ar.info(msg)
        else:
            count_statements_imported = num - len(failed_statements)
            msg = "Imported {0} statemenmts from file {1} with {2} statements failed.".format(count_statements_imported,
                                                                                filename,
                                                                                len(failed_statements))
            dd.logger.info(msg)
            for i, _statement in enumerate(failed_statements):
                msg_error = "Statement number {0} failed to be imported : {1}.".format(i + 1, failed_statements[_statement])
                dd.logger.info(msg_error)

dd.inject_action('system.SiteConfig', import_sepa=ImportStatements())

class Account(dd.Model):
    """A bank account related to a given :class:`Partner
    <lino.modlib.models.contacts.Partner>`.

    One partner can have more than one bank account.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Account')
        verbose_name = _("Account")
        verbose_name_plural = _("Accounts")

    partner = dd.ForeignKey(
        'contacts.Partner',
        related_name='sepa_accounts', null=True, blank=True)

    iban = IBANField(verbose_name=_("IBAN"))
    bic = BICField(verbose_name=_("BIC"), blank=True)

    remark = models.CharField(_("Remark"), max_length=200, blank=True)

    primary = models.BooleanField(
        _("Primary"),
        default=False,
        help_text=_(
            "Enabling this field will automatically disable any "
            "previous primary account and update "
            "the partner's IBAN and BIC"))

    allow_cascaded_delete = ['partner']

    def __unicode__(self):
        if self.remark:
            return "{0} ({1})".format(self.iban, self.remark)
        return self.iban

    def full_clean(self):
        if self.iban and not self.bic:
            if self.iban[0].isdigit():
                iban, bic = belgian_nban_to_iban_bic(self.iban)
                self.bic = bic
                self.iban = iban
            else:
                self.bic = iban2bic(self.iban) or ''
        super(Account, self).full_clean()

    def after_ui_save(self, ar, cw):
        super(Account, self).after_ui_save(ar, cw)
        if self.primary:
            mi = self.partner
            for o in mi.sepa_accounts.exclude(id=self.id):
                if o.primary:
                    o.primary = False
                    o.save()
                    ar.set_response(refresh_all=True)
            watcher = ChangeWatcher(mi)
            for k in PRIMARY_FIELDS:
                setattr(mi, k, getattr(self, k))
            mi.save()
            watcher.send_update(ar.request)

PRIMARY_FIELDS = dd.fields_list(Account, 'iban bic')

class Statement(dd.Model):
    """A bank statement.

    This data is automaticaly imported by :class:`ImportStatements`.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Statement')
        verbose_name = _("Statement")
        verbose_name_plural = _("Statements")

    def __unicode__(self):
        if self.account:
            if self.date:
                return "{0} ({1})".format(self.account, self.date)
            else:
                return self.account
        return ''

    account = dd.ForeignKey('sepa.Account')
    date = models.DateField(_('Date'), null=True)
    date_done = models.DateTimeField(_('Import Date'), null=True)
    statement_number = models.CharField(_('Statement number'), null=False, max_length=128)
    balance_start = dd.PriceField(_("Initial amount"), null=True)
    balance_end = dd.PriceField(_("Final amount"), null=True)
    balance_end_real = dd.PriceField(_("Real end balance"), null=True)
    currency_code = models.CharField(_('Currency'), max_length=3)

    # fields like statement_number, date, solde_initial, solde_final

class Movement(dd.Model):
    """A movement within a bank statement.

    This data is automaticaly imported by :class:`ImportStatements`.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Movement')
        verbose_name = _("Movement")
        verbose_name_plural = _("Movements")

    statement = dd.ForeignKey('sepa.Statement')
    unique_import_id = models.CharField(_('Unique import ID'), max_length=128)
    # movement_number = models.CharField(_("Ref of Mov"), null=False, max_length=32)
    movement_date = models.DateField(_('Movement date'), null=True)
    amount = dd.PriceField(_('Amount'), null=True)
    partner = models.ForeignKey('contacts.Partner', related_name='sepa_movement', null=True)
    partner_name = models.CharField(_('Partner name'), max_length=35)
    bank_account = dd.ForeignKey('sepa.Account', blank=True, null=True)
    ref = models.CharField(_('Ref'), null=False, max_length=35)

from .ui import *
<commit_msg>http://luc.lino-framework.org/blog/2015/1016.html<commit_after># -*- coding: UTF-8 -*-
# Copyright 2014-2015 Luc Saffre
# This file is part of Lino Cosi.
#
# Lino Cosi is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Lino Cosi is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with Lino Cosi.  If not, see
# <http://www.gnu.org/licenses/>.

"""
Database models for `lino_cosi.lib.sepa`.

"""

from __future__ import unicode_literals
import logging
from pprint import pformat
import glob
import os
from django.db import models
from django.core.exceptions import MultipleObjectsReturned
from lino.api import dd, _, rt
from lino.core.utils import ChangeWatcher
from .camt import CamtParser
from .fields import IBANField, BICField
from .utils import belgian_nban_to_iban_bic, iban2bic
import time

logger = logging.getLogger(__name__)

class ImportStatements(dd.Action):
    """Import the .xml files found in the directory specified at
    :attr:`import_statements_path
    <lino_cosi.lib.sepa.Plugin.import_statements_path>`.

    End-users invoke this via the menu command :menuselection:`SEPA
    --> Import SEPA`.

    When a file has been successfully imported, Lino deletes it.

    It might happen that an .xml file accidentally gets downloaded a
    second time. Lino does not create these statements again.

    """
    label = _("Import SEPA")
    http_method = 'POST'
    select_rows = False

    def get_view_permission(self, profile):
        """Make it invisible when :attr:`import_statements_path
        <lino_cosi.lib.sepa.Plugin.import_statements_path>` is empty.

        """
        if not dd.plugins.sepa.import_statements_path:
            return False
        return super(ImportStatements, self).get_view_permission(profile)

    def run_from_ui(self, ar):
        pth = dd.plugins.sepa.import_statements_path
        wc = os.path.join(pth, '*.xml')
        dd.logger.info("Importing SEPA statements from %s...", wc)
        count = 0
        if pth:
            for filename in glob.iglob(wc):
                self.import_file(ar, filename)
                count += 1
            msg = "{0} xml files would have been imported.".format(count)
            dd.logger.info(msg)
            return ar.success(msg, alert=_("Success"))
        msg = "No import_statements_path configured."
        return ar.error(msg, alert=_("Error"))

    def import_file(self, ar, filename):
        Account = rt.modules.sepa.Account

        # Parse a CAMT053 XML file.
        parser = CamtParser()
        data_file = open(filename, 'rb').read()
        num = 0
        failed_statements = {}
        try:
            dd.logger.info("Parsing %s with camt.", filename)
            res = parser.parse(data_file)
            if res is None:
                raise Exception("res is None")
            for _statement in res:
                iban = _statement['account_number']
                if iban is None:
                    # msg = "Statement without account number : {0}"
                    failed_statements[num] = "IBAN Not found"
                    continue
                    # raise Exception(msg.format(pformat(_statement)))
                try:
                    account = Account.objects.get(iban=iban)
                except Account.DoesNotExist:
                    account = Account(iban=iban)
                    account.full_clean()
                    account.save()
                except MultipleObjectsReturned:
                    msg = "Found more than one account with IBAN {0}"
                    dd.logger.warning(msg.format(iban))
                    continue
                s = Statement(account=account,
                              date=_statement['date'].strftime("%Y-%m-%d"),
                              date_done=time.strftime("%Y-%m-%d"),
                              statement_number=_statement['name'],
                              balance_end=_statement['balance_end'],
                              balance_start=_statement['balance_start'],
                              balance_end_real=_statement['balance_end_real'],
                              currency_code=_statement['currency_code'])
                s.save()
                for _movement in _statement['transactions']:
                    assert _movement.remote_account == account.iban
                    # TODO :check if the movement is already imported.
                    if Movement.objects.filter(
                        unique_import_id=_movement['unique_import_id']).exists():
                        continue
                    _ref = _movement.get('ref', '') or ''
                    m = Movement(statement=s,
                                 unique_import_id=_movement['unique_import_id'],
                                 movement_date=_movement['date'],
                                 amount=_movement['amount'],
                                 partner_name=_movement.remote_owner,
                                 ref=_ref,
                                 bank_account=account)
                    m.save()
                num += 1

        except ValueError:
            dd.logger.info("Statement file was not a camt file.")

        if len(failed_statements) == 0:
            msg = "Imported {0} statemenmts from file {1}.".format(num, filename)
            dd.logger.info(msg)
            ar.info(msg)
        else:
            count_statements_imported = num - len(failed_statements)
            msg = "Imported {0} statemenmts from file {1} with {2} statements failed.".format(count_statements_imported,
                                                                                filename,
                                                                                len(failed_statements))
            dd.logger.info(msg)
            for i, _statement in enumerate(failed_statements):
                msg_error = "Statement number {0} failed to be imported : {1}.".format(i + 1, failed_statements[_statement])
                dd.logger.info(msg_error)

dd.inject_action('system.SiteConfig', import_sepa=ImportStatements())

class Account(dd.Model):
    """A bank account related to a given :class:`Partner
    <lino.modlib.models.contacts.Partner>`.

    One partner can have more than one bank account.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Account')
        verbose_name = _("Account")
        verbose_name_plural = _("Accounts")

    partner = dd.ForeignKey(
        'contacts.Partner',
        related_name='sepa_accounts', null=True, blank=True)

    iban = IBANField(verbose_name=_("IBAN"))
    bic = BICField(verbose_name=_("BIC"), blank=True)

    remark = models.CharField(_("Remark"), max_length=200, blank=True)

    primary = models.BooleanField(
        _("Primary"),
        default=False,
        help_text=_(
            "Enabling this field will automatically disable any "
            "previous primary account and update "
            "the partner's IBAN and BIC"))

    allow_cascaded_delete = ['partner']

    def __unicode__(self):
        if self.remark:
            return "{0} ({1})".format(self.iban, self.remark)
        return self.iban

    def full_clean(self):
        if self.iban and not self.bic:
            if self.iban[0].isdigit():
                iban, bic = belgian_nban_to_iban_bic(self.iban)
                self.bic = bic
                self.iban = iban
            else:
                self.bic = iban2bic(self.iban) or ''
        super(Account, self).full_clean()

    def after_ui_save(self, ar, cw):
        super(Account, self).after_ui_save(ar, cw)
        if self.primary:
            mi = self.partner
            for o in mi.sepa_accounts.exclude(id=self.id):
                if o.primary:
                    o.primary = False
                    o.save()
                    ar.set_response(refresh_all=True)
            watcher = ChangeWatcher(mi)
            for k in PRIMARY_FIELDS:
                setattr(mi, k, getattr(self, k))
            mi.save()
            watcher.send_update(ar.request)

PRIMARY_FIELDS = dd.fields_list(Account, 'iban bic')

class Statement(dd.Model):
    """A bank statement.

    This data is automaticaly imported by :class:`ImportStatements`.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Statement')
        verbose_name = _("Statement")
        verbose_name_plural = _("Statements")

    def __unicode__(self):
        if self.account:
            if self.date:
                return "{0} ({1})".format(self.account, self.date)
            else:
                return self.account
        return ''

    account = dd.ForeignKey('sepa.Account')
    date = models.DateField(_('Date'), null=True)
    date_done = models.DateTimeField(_('Import Date'), null=True)
    statement_number = models.CharField(_('Statement number'), null=False, max_length=128)
    balance_start = dd.PriceField(_("Initial amount"), null=True)
    balance_end = dd.PriceField(_("Final amount"), null=True)
    balance_end_real = dd.PriceField(_("Real end balance"), null=True)
    currency_code = models.CharField(_('Currency'), max_length=3)

    # fields like statement_number, date, solde_initial, solde_final

class Movement(dd.Model):
    """A movement within a bank statement.

    This data is automaticaly imported by :class:`ImportStatements`.

    """

    class Meta:
        abstract = dd.is_abstract_model(__name__, 'Movement')
        verbose_name = _("Movement")
        verbose_name_plural = _("Movements")

    statement = dd.ForeignKey('sepa.Statement')
    unique_import_id = models.CharField(_('Unique import ID'), max_length=128)
    # movement_number = models.CharField(_("Ref of Mov"), null=False, max_length=32)
    movement_date = models.DateField(_('Movement date'), null=True)
    amount = dd.PriceField(_('Amount'), null=True)
    partner = models.ForeignKey('contacts.Partner', related_name='sepa_movement', null=True)
    partner_name = models.CharField(_('Partner name'), max_length=35)
    bank_account = dd.ForeignKey('sepa.Account', blank=True, null=True)
    ref = models.CharField(_('Ref'), null=False, max_length=35)

from .ui import *
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import random
import re

from util import hook

@hook.command
def munge(inp, munge_count=0):
    reps = 0
    for n in xrange(len(inp)):
        rep = character_replacements.get(inp[n])
        if rep:
            inp = inp[:n] + rep.decode('utf8') + inp[n + 1:]
            reps += 1
            if reps == munge_count:
                break
    return inp

def winnow(inputs, limit=400):
    "remove random elements from the list until it's short enough"
    combiner = lambda l: u', '.join(l)
    suffix = ''
    while len(combiner(inputs)) >= limit:
        inputs.pop(random.randint(0, len(inputs) - 1))
        suffix = ' ...'
    return combiner(inputs) + suffix

def add_tag(db, chan, nick, subject):
    match = db.execute('select * from tag where lower(nick)=lower(?) and'
                       ' chan=? and lower(subject)=lower(?)',
                       (nick, chan, subject)).fetchall()
    if match:
        return 'already tagged'

    db.execute('replace into tag(chan, subject, nick) values(?,?,?)',
               (chan, subject, nick))
    db.commit()

    return 'tag added'

def delete_tag(db, chan, nick, del_tag):
    count = db.execute('delete from tag where lower(nick)=lower(?) and'
                       ' chan=? and lower(subject)=lower(?)',
                       (nick, chan, del_tag)).rowcount
    db.commit()

    if count:
        return 'deleted'
    else:
        return 'tag not found'

def get_tag_counts_by_chan(db, chan):
    tags = db.execute("select subject, count(*) from tag where chan=?"
                      " group by lower(subject)"
                      " order by lower(subject)", (chan,)).fetchall()

    tags.sort(key=lambda x: x[1], reverse=True)
    if not tags:
        return 'no tags in %s' % chan
    ret = '%s tags: ' % chan
    return winnow(['%s (%d)' % row for row in tags])

def get_tags_by_nick(db, chan, nick):
    return db.execute("select subject from tag where lower(nick)=lower(?)"
                      " and chan=?"
                      " order by lower(subject)", (nick, chan)).fetchall()

def get_nicks_by_tagset(db, chan, tagset):
    nicks = None
    for tag in tagset.split('&'):
        tag = tag.strip()

        current_nicks = db.execute("select nick from tag where " +
                                   "lower(subject)=lower(?)"
                                   " and chan=?", (tag, chan)).fetchall()

        if not current_nicks:
            return "tag '%s' not found" % tag

        if nicks is None:
            nicks = set(current_nicks)
        else:
            nicks.intersection_update(current_nicks)

    nicks = [munge(x[0], 1) for x in sorted(nicks)]
    if not nicks:
        return 'no tags found in intersection of "%s"' % tagset
    return 'nicks tagged "%s": ' % tagset + winnow(nicks)

@hook.command
def tag(inp, chan='', db=None):
    '.tag <nick>/[add|del] <nick> <tag>/list [tag] [& tag...] -- get list of tags on ' \
    '<nick>/(un)marks <nick> as <tag>/gets list of tags/nicks marked as [tag]'

    db.execute('create table if not exists tag(chan, subject, nick)')

    add = re.match(r'(?:a(?:dd)? )?(\S+) (.+)', inp)
    delete = re.match(r'd(?:el(?:ete)?)? (\S+) (.+)\s*$', inp)
    retrieve = re.match(r'l(?:ist)(?: (.+))?\s*$', inp)

    if retrieve:
        search_tag = retrieve.group(1)
        if search_tag:
            return get_nicks_by_tagset(db, chan, search_tag)
        else:
            return get_tag_counts_by_chan(db, chan)
    if delete:
        nick, del_tag = delete.groups()
        return delete_tag(db, chan, nick, del_tag)
    if add:
        nick, subject = add.groups()
        return add_tag(db, chan, nick, subject)
    else:
        tags = get_tags_by_nick(db, chan, inp)

        if not tags:
            return get_nicks_by_tagset(db, chan, inp)
        else:
            return 'tags for "%s": ' % munge(inp, 1) + winnow([
                tag[0] for tag in tags])

character_replacements = {
    'a': 'ä',
#    'b': 'Б',
    'c': 'ċ',
    'd': 'đ',
    'e': 'ë',
    'f': 'ƒ',
    'g': 'ġ',
    'h': 'ħ',
    'i': 'í',
    'j': 'ĵ',
    'k': 'ķ',
    'l': 'ĺ',
#    'm': 'ṁ',
    'n': 'ñ',
    'o': 'ö',
    'p': 'ρ',
#    'q': 'ʠ',
    'r': 'ŗ',
    's': 'š',
    't': 'ţ',
    'u': 'ü',
#    'v': '',
    'w': 'ω',
    'x': 'χ',
    'y': 'ÿ',
    'z': 'ź',
    'A': 'Å',
    'B': 'Β',
    'C': 'Ç',
    'D': 'Ď',
    'E': 'Ē',
#    'F': 'Ḟ',
    'G': 'Ġ',
    'H': 'Ħ',
    'I': 'Í',
    'J': 'Ĵ',
    'K': 'Ķ',
    'L': 'Ĺ',
    'M': 'Μ',
    'N': 'Ν',
    'O': 'Ö',
    'P': 'Р',
#    'Q': 'Ｑ',
    'R': 'Ŗ',
    'S': 'Š',
    'T': 'Ţ',
    'U': 'Ů',
#    'V': 'Ṿ',
    'W': 'Ŵ',
    'X': 'Χ',
    'Y': 'Ỳ',
    'Z': 'Ż'}
<commit_msg>tag.py: split .tag into .tag/.tags/.tagged, add automatic pagination<commit_after># -*- coding: utf-8 -*-

import random
import re
import threading

from util import hook

@hook.command
def munge(inp, munge_count=0):
    reps = 0
    for n in xrange(len(inp)):
        rep = character_replacements.get(inp[n])
        if rep:
            inp = inp[:n] + rep.decode('utf8') + inp[n + 1:]
            reps += 1
            if reps == munge_count:
                break
    return inp

class PaginatingWinnower(object):
    def __init__(self):
        self.lock = threading.Lock()
        self.last_input = []
        self.recent = set()

    def winnow(self, inputs, limit=400, ordered=False):
        "remove random elements from the list until it's short enough"
        with self.lock:
            # try to remove elements that were *not* removed recently
            inputs_sorted = sorted(inputs)
            if inputs_sorted == self.last_input:
                same_input = True
            else:
                same_input = False
                self.last_input = inputs_sorted
                self.recent.clear()

            combiner = lambda l: u', '.join(l)
            suffix = ''

            while len(combiner(inputs)) >= limit:
                if same_input and any(inp in self.recent for inp in inputs):
                    if ordered:
                        for inp in self.recent:
                            if inp in inputs:
                                inputs.remove(inp)
                    else:
                        inputs.remove(random.choice([inp for inp in inputs if inp in self.recent]))
                else:
                    if ordered:
                        inputs.pop()
                    else:
                        inputs.pop(random.randint(0, len(inputs) - 1))
                suffix = ' ...'

            self.recent.update(inputs)
            return combiner(inputs) + suffix

winnow = PaginatingWinnower().winnow

def add_tag(db, chan, nick, subject):
    match = db.execute('select * from tag where lower(nick)=lower(?) and'
                       ' chan=? and lower(subject)=lower(?)',
                       (nick, chan, subject)).fetchall()
    if match:
        return 'already tagged'

    db.execute('replace into tag(chan, subject, nick) values(?,?,?)',
               (chan, subject, nick))
    db.commit()

    return 'tag added'

def delete_tag(db, chan, nick, del_tag):
    count = db.execute('delete from tag where lower(nick)=lower(?) and'
                       ' chan=? and lower(subject)=lower(?)',
                       (nick, chan, del_tag)).rowcount
    db.commit()

    if count:
        return 'deleted'
    else:
        return 'tag not found'

def get_tag_counts_by_chan(db, chan):
    tags = db.execute("select subject, count(*) from tag where chan=?"
                      " group by lower(subject)"
                      " order by lower(subject)", (chan,)).fetchall()

    tags.sort(key=lambda x: x[1], reverse=True)
    if not tags:
        return 'no tags in %s' % chan
    ret = '%s tags: ' % chan
    return winnow(['%s (%d)' % row for row in tags], ordered=True)

def get_tags_by_nick(db, chan, nick):
    tags = db.execute("select subject from tag where lower(nick)=lower(?)"
                      " and chan=?"
                      " order by lower(subject)", (nick, chan)).fetchall()
    if tags:
        return 'tags for "%s": ' % munge(nick, 1) + winnow([
                tag[0] for tag in tags])
    else:
        return ''

def get_nicks_by_tagset(db, chan, tagset):
    nicks = None
    for tag in tagset.split('&'):
        tag = tag.strip()

        current_nicks = db.execute("select nick from tag where " +
                                   "lower(subject)=lower(?)"
                                   " and chan=?", (tag, chan)).fetchall()

        if not current_nicks:
            return "tag '%s' not found" % tag

        if nicks is None:
            nicks = set(current_nicks)
        else:
            nicks.intersection_update(current_nicks)

    nicks = [munge(x[0], 1) for x in sorted(nicks)]
    if not nicks:
        return 'no tags found in intersection of "%s"' % tagset
    return 'nicks tagged "%s": ' % tagset + winnow(nicks)

@hook.command
def tag(inp, chan='', db=None):
    '.tag [add|del] <nick> <tag> -- marks/unmarks <nick> as <tag> {related: .tags, .tagged}'

    db.execute('create table if not exists tag(chan, subject, nick)')

    add = re.match(r'(?:a(?:dd)? )?(\S+) (.+)', inp)
    delete = re.match(r'd(?:el(?:ete)?)? (\S+) (.+)\s*$', inp)

    if delete:
        nick, del_tag = delete.groups()
        return delete_tag(db, chan, nick, del_tag)
    if add:
        nick, subject = add.groups()
        return add_tag(db, chan, nick, subject)
    else:
        tags = get_tags_by_nick(db, chan, inp)
        if tags:
            return tags
        else:
            return tag.__doc__

@hook.command
def tags(inp, chan='', db=None):
    '.tags <nick>/list -- get list of tags for <nick>, or a list of tags {related: .tag, .tagged}'
    if inp == 'list':
        return get_tag_counts_by_chan(db, chan)

    tags = get_tags_by_nick(db, chan, inp)
    if tags:
        return tags
    else:
        return get_nicks_by_tagset(db, chan, inp)

@hook.command
def tagged(inp, chan='', db=None):
    '.tagged <tag> [& tag...] -- get nicks marked as <tag> (separate multiple tags with &) {related: .tag, .tags}'

    return get_nicks_by_tagset(db, chan, inp)

character_replacements = {
    'a': 'ä',
#    'b': 'Б',
    'c': 'ċ',
    'd': 'đ',
    'e': 'ë',
    'f': 'ƒ',
    'g': 'ġ',
    'h': 'ħ',
    'i': 'í',
    'j': 'ĵ',
    'k': 'ķ',
    'l': 'ĺ',
#    'm': 'ṁ',
    'n': 'ñ',
    'o': 'ö',
    'p': 'ρ',
#    'q': 'ʠ',
    'r': 'ŗ',
    's': 'š',
    't': 'ţ',
    'u': 'ü',
#    'v': '',
    'w': 'ω',
    'x': 'χ',
    'y': 'ÿ',
    'z': 'ź',
    'A': 'Å',
    'B': 'Β',
    'C': 'Ç',
    'D': 'Ď',
    'E': 'Ē',
#    'F': 'Ḟ',
    'G': 'Ġ',
    'H': 'Ħ',
    'I': 'Í',
    'J': 'Ĵ',
    'K': 'Ķ',
    'L': 'Ĺ',
    'M': 'Μ',
    'N': 'Ν',
    'O': 'Ö',
    'P': 'Р',
#    'Q': 'Ｑ',
    'R': 'Ŗ',
    'S': 'Š',
    'T': 'Ţ',
    'U': 'Ů',
#    'V': 'Ṿ',
    'W': 'Ŵ',
    'X': 'Χ',
    'Y': 'Ỳ',
    'Z': 'Ż'}
<|endoftext|>


<commit_before># vim: fileencoding=utf-8:
'''

Format
---------------------

Sources:

.. _Swift for corporates: http://www.sepaforcorporates.com/\
    swift-for-corporates/account-statement-mt940-file-format-overview/
.. _Rabobank MT940: https://www.rabobank.nl/images/\
    formaatbeschrijving_swift_bt940s_1_0_nl_rib_29539296.pdf

 - `Swift for corporates`_
 - `Rabobank MT940`_

::

    [] = optional
    ! = fixed length
    a = Text
    x = Alphanumeric, seems more like text actually. Can include special
        characters (slashes) and whitespace as well as letters and numbers
    d = Numeric separated by decimal (usually comma)
    c = Code list value
    n = Numeric
'''

import os

import mt940

def parse(src, encoding=None):
    '''
    Parses mt940 data and returns transactions object

    :param src: file handler to read, filename to read or raw data as string
    :return: Collection of transactions
    :rtype: Transactions
    '''

    if hasattr(src, 'read'):  # pragma: no branch
        data = src.read()
    elif os.path.isfile(src):
        return parse(open(src, 'rb').read())
    elif hasattr(src, 'decode'):
        exception = None
        encodings = [encoding, 'utf-8', 'cp852', 'iso8859-15', 'latin1']

        for encoding in encodings:  # pragma: no branch
            if not encoding:
                continue

            try:
                data = src.decode(encoding)
                break
            except UnicodeDecodeError as e:
                exception = e
        else:  # pragma: no cover
            raise exception
    else:  # pragma: no cover
        data = src

    transactions = mt940.models.Transactions()
    transactions.parse(data)

    return transactions
<commit_msg>fixed windows bug when directly passing data. Fixes #68<commit_after># vim: fileencoding=utf-8:
'''

Format
---------------------

Sources:

.. _Swift for corporates: http://www.sepaforcorporates.com/\
    swift-for-corporates/account-statement-mt940-file-format-overview/
.. _Rabobank MT940: https://www.rabobank.nl/images/\
    formaatbeschrijving_swift_bt940s_1_0_nl_rib_29539296.pdf

 - `Swift for corporates`_
 - `Rabobank MT940`_

::

    [] = optional
    ! = fixed length
    a = Text
    x = Alphanumeric, seems more like text actually. Can include special
        characters (slashes) and whitespace as well as letters and numbers
    d = Numeric separated by decimal (usually comma)
    c = Code list value
    n = Numeric
'''

import os

import mt940

def parse(src, encoding=None):
    '''
    Parses mt940 data and returns transactions object

    :param src: file handler to read, filename to read or raw data as string
    :return: Collection of transactions
    :rtype: Transactions
    '''

    def safe_is_file(filename):
        try:
            return os.path.isfile(src)
        except ValueError:  # pragma: no cover
            return False

    if hasattr(src, 'read'):  # pragma: no branch
        data = src.read()
    elif safe_is_file(src):
        return parse(open(src, 'rb').read())
    elif hasattr(src, 'decode'):
        exception = None
        encodings = [encoding, 'utf-8', 'cp852', 'iso8859-15', 'latin1']

        for encoding in encodings:  # pragma: no branch
            if not encoding:
                continue

            try:
                data = src.decode(encoding)
                break
            except UnicodeDecodeError as e:
                exception = e
        else:  # pragma: no cover
            raise exception
    else:  # pragma: no cover
        data = src

    transactions = mt940.models.Transactions()
    transactions.parse(data)

    return transactions
<|endoftext|>


<commit_before># coding=utf-8

import html.parser
import re
import urllib.request
import urllib.error

DEFAULT_AGENT = 'Yui'

class TitleParser(html.parser.HTMLParser):
    def __init__(self):
        html.parser.HTMLParser.__init__(self)
        self.reading = False
        self.done = False
        self.title = ''

    def handle_starttag(self, tag, attrs):
        if tag == 'title' and not self.done:
            self.reading = True

    def handle_endtag(self, tag):
        if tag == 'title':
            self.reading = False
            self.done = True

    def handle_data(self, data):
        if self.reading:
            self.title += data

    def handle_charref(self, ref):
        if self.reading:
            self.handle_entityref('#' + ref)

    def handle_entityref(self, ref):
        if self.reading:
            self.title += '&%s;' % ref

def url_encode_ascii(string):
    return re.sub('[\x80-\xFF]', lambda c: '%%%02x' % ord(c.group(0)), string)

def humanify(num):
    pre = ['KiB', 'MiB', 'GiB', 'TiB']
    if num < 1024:
        return '%d Byte' % num
    num /= 1024
    for p in pre:
        div = num / 1024
        if div < 1.0:
            return '%.2f%s' % (num, p)
        num = div

# returns a properly encoded url (escaped unicode etc)
# or None if it's not a valid url
def get_encoded_url(url):
    # test if it's a valid URL and encode it properly, if it is
    parts = urllib.request.urlparse(url)
    if not ((parts[0] == 'http' or parts[0] == 'https') and parts[1] and parts[1] != 'localhost' and not
    parts[1].split('.')[-1].isdigit()):
        return None

    # handle unicode URLs
    url = urllib.request.urlunparse(
        p if i == 1 else url_encode_ascii(p)
        for i, p in enumerate(parts)
    )
    return url

def get_url_title(url):
    enc = ['utf8', 'iso-8869-1', 'shift-jis']
    title = ''
    headers = {
        'User-Agent': yui.config_val('httpUserAgent', default=DEFAULT_AGENT)
    }
    try:
        req = urllib.request.Request(url, data=None, headers=headers)
        resp = urllib.request.urlopen(req, timeout=5)
    except urllib.error.HTTPError as e:
        return 'Status ' + e.code
    except urllib.error.URLError as e:
        return 'Error: ' + e.reason

    # try the charset set in the html header, if there is one
    if 'content-type' in resp.headers and 'charset=' in resp.headers['content-type']:
        enc = enc + [resp.headers['content-type'].split('charset=')[-1]]

    # read up to 1mb
    chunk = resp.read(1024 * 1024)
    parser = TitleParser()
    for e in enc:
        try:
            decoded_chunk = chunk.decode(e, 'ignore')
            parser.feed(decoded_chunk)
            if parser.done:
                title = parser.title
            parser.close()
            if len(title) > 0:
                esc = parser.unescape(title)
                return 'Title: '+esc.strip()
        except Exception as ex:
            pass

    # no title, try to output some other useful data
    info = []
    if 'content-type' in resp.headers:
        info.append('Type: ' + resp.headers['content-type'].split(';')[0])
    if 'content-length' in resp.headers:
        info.append('Size: ' + humanify(int(resp.headers['content-length'])))

    return ', '.join(info)

@yui.event('msgRecv')
def url(msg, channel):
    # find urls in channel message
    words = msg.split()
    titles = []
    max_urls = 3
    for w in words:
        enc_url = get_encoded_url(w)
        if not enc_url:
            continue

        title = get_url_title(enc_url)
        if title:
            titles.append(title.strip())

        if len(titles) >= max_urls:
            break

    # don't say anything, if we couldn't get any titles
    if len(titles) > 0:
        concat = ' \x035|\x03 '.join(titles)
        yui.send_msg(channel, concat)
<commit_msg>fix implicit int to string problem<commit_after># coding=utf-8

import html.parser
import re
import urllib.request
import urllib.error

DEFAULT_AGENT = 'Yui'

class TitleParser(html.parser.HTMLParser):
    def __init__(self):
        html.parser.HTMLParser.__init__(self)
        self.reading = False
        self.done = False
        self.title = ''

    def handle_starttag(self, tag, attrs):
        if tag == 'title' and not self.done:
            self.reading = True

    def handle_endtag(self, tag):
        if tag == 'title':
            self.reading = False
            self.done = True

    def handle_data(self, data):
        if self.reading:
            self.title += data

    def handle_charref(self, ref):
        if self.reading:
            self.handle_entityref('#' + ref)

    def handle_entityref(self, ref):
        if self.reading:
            self.title += '&%s;' % ref

def url_encode_ascii(string):
    return re.sub('[\x80-\xFF]', lambda c: '%%%02x' % ord(c.group(0)), string)

def humanify(num):
    pre = ['KiB', 'MiB', 'GiB', 'TiB']
    if num < 1024:
        return '%d Byte' % num
    num /= 1024
    for p in pre:
        div = num / 1024
        if div < 1.0:
            return '%.2f%s' % (num, p)
        num = div

# returns a properly encoded url (escaped unicode etc)
# or None if it's not a valid url
def get_encoded_url(url):
    # test if it's a valid URL and encode it properly, if it is
    parts = urllib.request.urlparse(url)
    if not ((parts[0] == 'http' or parts[0] == 'https') and parts[1] and parts[1] != 'localhost' and not
    parts[1].split('.')[-1].isdigit()):
        return None

    # handle unicode URLs
    url = urllib.request.urlunparse(
        p if i == 1 else url_encode_ascii(p)
        for i, p in enumerate(parts)
    )
    return url

def get_url_title(url):
    enc = ['utf8', 'iso-8869-1', 'shift-jis']
    title = ''
    headers = {
        'User-Agent': yui.config_val('httpUserAgent', default=DEFAULT_AGENT)
    }
    try:
        req = urllib.request.Request(url, data=None, headers=headers)
        resp = urllib.request.urlopen(req, timeout=5)
    except urllib.error.HTTPError as e:
        return 'Status ' + str(e.code)
    except urllib.error.URLError as e:
        return 'Error: ' + e.reason

    # try the charset set in the html header, if there is one
    if 'content-type' in resp.headers and 'charset=' in resp.headers['content-type']:
        enc = enc + [resp.headers['content-type'].split('charset=')[-1]]

    # read up to 1mb
    chunk = resp.read(1024 * 1024)
    parser = TitleParser()
    for e in enc:
        try:
            decoded_chunk = chunk.decode(e, 'ignore')
            parser.feed(decoded_chunk)
            if parser.done:
                title = parser.title
            parser.close()
            if len(title) > 0:
                esc = parser.unescape(title)
                return 'Title: '+esc.strip()
        except Exception as ex:
            pass

    # no title, try to output some other useful data
    info = []
    if 'content-type' in resp.headers:
        info.append('Type: ' + resp.headers['content-type'].split(';')[0])
    if 'content-length' in resp.headers:
        info.append('Size: ' + humanify(int(resp.headers['content-length'])))

    return ', '.join(info)

@yui.event('msgRecv')
def url(msg, channel):
    # find urls in channel message
    words = msg.split()
    titles = []
    max_urls = 3
    for w in words:
        enc_url = get_encoded_url(w)
        if not enc_url:
            continue

        title = get_url_title(enc_url)
        if title:
            titles.append(title.strip())

        if len(titles) >= max_urls:
            break

    # don't say anything, if we couldn't get any titles
    if len(titles) > 0:
        concat = ' \x035|\x03 '.join(titles)
        yui.send_msg(channel, concat)
<|endoftext|>


<commit_before><commit_msg>Revert 39618<commit_after><|endoftext|>


<commit_before>from small_multiples import small_multiples_plot

import os
import gc

from datetime import datetime, timedelta
import numpy as np
import pupynere as nc

# from pylab import figure, get_cmap, colorbar
from matplotlib.figure import figaspect, Figure
from matplotlib.colorbar import ColorbarBase
from matplotlib.cm import get_cmap
from matplotlib.ticker import FuncFormatter
from matplotlib.dates import mx2num, date2num, DateFormatter
from matplotlib.backends.backend_agg import FigureCanvasAgg

from math import ceil

import pytz
tz=pytz.timezone('US/Eastern') # Why, oh, why, is it using my local time zone?
time_series_x_fmt = DateFormatter('%H%M:%S', tz=tz)

def kilo(x, pos):
    'The two args are the value and tick position'
    return '%s' % (x/1000.0)

kilo_formatter = FuncFormatter(kilo)

# # /* DC LMA */
# DClat = 38.8888500 # falls church / western tip of arlington, rough centroid of stations
# DClon = -77.1685800
# kounLat = 35.23833
# kounLon = -97.46028
# kounAlt = 377.0
# radarLat=kounLat
# radarLon=kounLon
# radarAlt=0.0

# # ARMOR
# radarLat = 34.6461
# radarLon = -86.7714
# radarAlt = 180.0

# mapProj = MapProjection(projection='eqc', ctrLat=radarLat, ctrLon=radarLon, lat_ts=radarLat, lon_0=radarLon)
# geoProj = GeographicSystem()

def centers_to_edges(x):
    xedge=np.zeros(x.shape[0]+1)
    xedge[1:-1] = (x[:-1] + x[1:])/2.0
    dx = np.mean(np.abs(xedge[2:-1] - xedge[1:-2]))
    xedge[0] = xedge[1] - dx
    xedge[-1] = xedge[-2] + dx
    return xedge

def multiples_figaspect(n_rows, n_cols, x_range, y_range, fig_width=None, max_height=None):
    """
        DIMENSIONS, assuming no margins anywhere.
        Plottable area: w, h (inches)
        Data units: x_range, y_range (km, array shape, etc.)
        Subaxis count: n_cols, n_rows (laid out in a uniformly spaced grid)
        Subaxis size: x_box, y_box (inches)

        y_box = h / n_row
        x_box = w / n_col

        Assume we're working with a plottable area of fig_width (no margins).
        Also assume we want squared-up units (e.g, square
        grid pixels if x_range, y_range are nx and ny for an array, or physical
        units, so 1 km north = 1 km east). In that case, we want:

            y_range / y_box  =  x_range / x_box  (e.g., km/in)

        So, given w = fig_width (taken from rcParams by default):

            (y_range * n_rows) / h  =  (x_range * n_col) / w

            or

            h = (y_range * n_rows * w) / (x_range * n_col)

        If we want to calculate the number of rows that can fit on a page with max_height,

            y_box = h / n_rows
            n_rows_perpage = floor(max_height / y_box)

        Returns: w, h, n_rows_perpage, n_pages

    """
    if fig_width is None:
        fig_width, fig_height = matplotlib.rcParams['figure.figsize'][0:1]
    w = float(fig_width)

    n_rows, n_cols, x_range, y_range = (float(i) for i in (n_rows, n_cols, x_range, y_range))

    h = (y_range * n_rows * w) / (x_range * n_cols)

    if max_height is not None:
        max_height = float(max_height)
        y_box = h / n_row
        n_rows_perpage = floor(max_height / y_box)
        n_pages = ceil(n_rows/n_rows_perpage)
    else:
        n_rows_perpage = n_rows
        n_pages = 1

    return w, h, n_rows_perpage, n_pages

def make_plot(filename, grid_name, x_name='x', y_name='y', t_name='time',
                n_cols=6, outpath='', filename_prefix='LMA',
                do_save=True, image_type='pdf'):

    f = nc.NetCDFFile(filename)
    data = f.variables  # dictionary of variable names to nc_var objects
    dims = f.dimensions # dictionary of dimension names to sizes
    x = data[x_name]
    y = data[y_name]
    t = data[t_name]
    grid = data[grid_name]

    assert len(x.shape) == 1
    assert len(y.shape) == 1
    assert len(t.shape) == 1

    grid_dims = grid.dimensions # tuple of dimension names
    name_to_idx = dict((k, i) for i, k in enumerate(grid_dims))

    grid_t_idx = name_to_idx[t.dimensions[0]]
    grid_x_idx = name_to_idx[x.dimensions[0]]
    grid_y_idx = name_to_idx[y.dimensions[0]]

    n_frames = t.shape[0]
    # n_cols = 6
    n_rows = int(ceil( float(n_frames) / n_cols ))

    colormap = get_cmap('gist_earth')
    grey_color = (0.5,)*3
    frame_color = (0.2,)*3

    density_maxes = []
    total_counts = []
    all_t = []

    xedge = centers_to_edges(x)
    x_range = xedge.max() - xedge.min()
    yedge = centers_to_edges(y)
    y_range = yedge.max() - yedge.min()
    dx = (xedge[1]-xedge[0])

    # w, h = figaspect(float(n_rows)/n_cols) # breaks for large numbers of frames - has a hard-coded max figure size
    w, h, n_rows_perpage, n_pages = multiples_figaspect(n_rows, n_cols, x_range, y_range, fig_width=8.5, max_height=None)

    # count_scale_factor = dx # / 1000.0
    # max_count_baseline = 450 * count_scale_factor #/ 10.0
    min_count, max_count = 1, grid[:].max() #max_count_baseline*(t[1]-t[0])

    f.close()

    fig = Figure(figsize=(w,h))
    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)
    p = small_multiples_plot(fig=fig, rows=n_rows, columns=n_cols)
    p.label_edges(True)
    pad = 0.0 # for time labels in each frame

    for ax in p.multiples.flat:
        ax.set_axis_bgcolor('black')
        ax.spines['top'].set_edgecolor(frame_color)
        ax.spines['bottom'].set_edgecolor(frame_color)
        ax.spines['left'].set_edgecolor(frame_color)
        ax.spines['right'].set_edgecolor(frame_color)
    #     ax.yaxis.set_major_formatter(kilo_formatter)
    #     ax.xaxis.set_major_formatter(kilo_formatter)
    base_date = datetime.strptime(t.units, "seconds since %Y-%m-%d %H:%M:%S")
    time_delta = timedelta(0,float(t[1]-t[0]),0)
    start_time = base_date + time_delta

    indexer = [None,]*len(grid.shape)

    frame_start_times = []
    for i in range(n_frames):
        frame_start = base_date + timedelta(0,float(t[i]),0)
        frame_start_times.append(frame_start)
        indexer[grid_t_idx] = i

        density = grid[indexer]

        # density,edges = np.histogramdd((x,y), bins=(xedge,yedge))
        density_plot  = p.multiples.flat[i].pcolormesh(xedge,yedge,
                                   np.log10(density.transpose()),
                                   vmin=-0.2,
                                   vmax=np.log10(max_count),
                                   cmap=colormap)
        label_string = frame_start.strftime('%H%M:%S')
        text_label = p.multiples.flat[i].text(xedge[0]-pad+x_range*.015, yedge[0]-pad+y_range*.015, label_string, color=grey_color, size=6)
        density_plot.set_rasterized(True)
        density_maxes.append(density.max())
        total_counts.append(density.sum())
        all_t.append(frame_start)
        print label_string, x.shape, density.max(), density.sum()

    color_scale = ColorbarBase(p.colorbar_ax, cmap=density_plot.cmap,
                                       norm=density_plot.norm,
                                       orientation='horizontal')

    # color_scale.set_label('count per pixel')
    color_scale.set_label('log10(count per pixel)')

    view_x = (xedge.min(), xedge.max())
    view_y = (yedge.min(), yedge.max())

    print 'making multiples',
    p.multiples.flat[0].axis(view_x+view_y)
    filename = '%s-%s_%s_%05.2fkm_%05.1fs.%s' % (filename_prefix, grid_name, start_time.strftime('%Y%m%d_%H%M%S'), dx, time_delta.seconds, image_type)
    filename = os.path.join(outpath, filename)
    if do_save:
        fig.savefig(filename, dpi=150)

    return fig, p, frame_start_times, filename

    print ' ... done'

if __name__ == '__main__':
    do_profile=False
    if do_profile:
        import hotshot
        from hotshot import stats
        prof = hotshot.Profile("multiples_test_profile")
        prof.runcall(runtest)
        prof.close()
        s=stats.load("multiples_test_profile")
        s.sort_stats("time").print_stats()
    else:
        import sys
        res = runtest(sys.argv[1], sys.argv[2])<commit_msg>Assign sane max value for colorbar if grid is <= 1 everywhere<commit_after>from small_multiples import small_multiples_plot

import os
import gc

from datetime import datetime, timedelta
import numpy as np
import pupynere as nc

# from pylab import figure, get_cmap, colorbar
from matplotlib.figure import figaspect, Figure
from matplotlib.colorbar import ColorbarBase
from matplotlib.cm import get_cmap
from matplotlib.ticker import FuncFormatter
from matplotlib.dates import mx2num, date2num, DateFormatter
from matplotlib.backends.backend_agg import FigureCanvasAgg

from math import ceil

import pytz
tz=pytz.timezone('US/Eastern') # Why, oh, why, is it using my local time zone?
time_series_x_fmt = DateFormatter('%H%M:%S', tz=tz)

def kilo(x, pos):
    'The two args are the value and tick position'
    return '%s' % (x/1000.0)

kilo_formatter = FuncFormatter(kilo)

# # /* DC LMA */
# DClat = 38.8888500 # falls church / western tip of arlington, rough centroid of stations
# DClon = -77.1685800
# kounLat = 35.23833
# kounLon = -97.46028
# kounAlt = 377.0
# radarLat=kounLat
# radarLon=kounLon
# radarAlt=0.0

# # ARMOR
# radarLat = 34.6461
# radarLon = -86.7714
# radarAlt = 180.0

# mapProj = MapProjection(projection='eqc', ctrLat=radarLat, ctrLon=radarLon, lat_ts=radarLat, lon_0=radarLon)
# geoProj = GeographicSystem()

def centers_to_edges(x):
    xedge=np.zeros(x.shape[0]+1)
    xedge[1:-1] = (x[:-1] + x[1:])/2.0
    dx = np.mean(np.abs(xedge[2:-1] - xedge[1:-2]))
    xedge[0] = xedge[1] - dx
    xedge[-1] = xedge[-2] + dx
    return xedge

def multiples_figaspect(n_rows, n_cols, x_range, y_range, fig_width=None, max_height=None):
    """
        DIMENSIONS, assuming no margins anywhere.
        Plottable area: w, h (inches)
        Data units: x_range, y_range (km, array shape, etc.)
        Subaxis count: n_cols, n_rows (laid out in a uniformly spaced grid)
        Subaxis size: x_box, y_box (inches)

        y_box = h / n_row
        x_box = w / n_col

        Assume we're working with a plottable area of fig_width (no margins).
        Also assume we want squared-up units (e.g, square
        grid pixels if x_range, y_range are nx and ny for an array, or physical
        units, so 1 km north = 1 km east). In that case, we want:

            y_range / y_box  =  x_range / x_box  (e.g., km/in)

        So, given w = fig_width (taken from rcParams by default):

            (y_range * n_rows) / h  =  (x_range * n_col) / w

            or

            h = (y_range * n_rows * w) / (x_range * n_col)

        If we want to calculate the number of rows that can fit on a page with max_height,

            y_box = h / n_rows
            n_rows_perpage = floor(max_height / y_box)

        Returns: w, h, n_rows_perpage, n_pages

    """
    if fig_width is None:
        fig_width, fig_height = matplotlib.rcParams['figure.figsize'][0:1]
    w = float(fig_width)

    n_rows, n_cols, x_range, y_range = (float(i) for i in (n_rows, n_cols, x_range, y_range))

    h = (y_range * n_rows * w) / (x_range * n_cols)

    if max_height is not None:
        max_height = float(max_height)
        y_box = h / n_row
        n_rows_perpage = floor(max_height / y_box)
        n_pages = ceil(n_rows/n_rows_perpage)
    else:
        n_rows_perpage = n_rows
        n_pages = 1

    return w, h, n_rows_perpage, n_pages

def make_plot(filename, grid_name, x_name='x', y_name='y', t_name='time',
                n_cols=6, outpath='', filename_prefix='LMA',
                do_save=True, image_type='pdf'):

    f = nc.NetCDFFile(filename)
    data = f.variables  # dictionary of variable names to nc_var objects
    dims = f.dimensions # dictionary of dimension names to sizes
    x = data[x_name]
    y = data[y_name]
    t = data[t_name]
    grid = data[grid_name]

    assert len(x.shape) == 1
    assert len(y.shape) == 1
    assert len(t.shape) == 1

    grid_dims = grid.dimensions # tuple of dimension names
    name_to_idx = dict((k, i) for i, k in enumerate(grid_dims))

    grid_t_idx = name_to_idx[t.dimensions[0]]
    grid_x_idx = name_to_idx[x.dimensions[0]]
    grid_y_idx = name_to_idx[y.dimensions[0]]

    n_frames = t.shape[0]
    # n_cols = 6
    n_rows = int(ceil( float(n_frames) / n_cols ))

    colormap = get_cmap('gist_earth')
    grey_color = (0.5,)*3
    frame_color = (0.2,)*3

    density_maxes = []
    total_counts = []
    all_t = []

    xedge = centers_to_edges(x)
    x_range = xedge.max() - xedge.min()
    yedge = centers_to_edges(y)
    y_range = yedge.max() - yedge.min()
    dx = (xedge[1]-xedge[0])

    # w, h = figaspect(float(n_rows)/n_cols) # breaks for large numbers of frames - has a hard-coded max figure size
    w, h, n_rows_perpage, n_pages = multiples_figaspect(n_rows, n_cols, x_range, y_range, fig_width=8.5, max_height=None)

    # count_scale_factor = dx # / 1000.0
    # max_count_baseline = 450 * count_scale_factor #/ 10.0
    min_count, max_count = 1, grid[:].max() #max_count_baseline*(t[1]-t[0])
    if (max_count == 0) | (max_count == 1 ):
        max_count = min_count+1
    f.close()

    fig = Figure(figsize=(w,h))
    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)
    p = small_multiples_plot(fig=fig, rows=n_rows, columns=n_cols)
    p.label_edges(True)
    pad = 0.0 # for time labels in each frame

    for ax in p.multiples.flat:
        ax.set_axis_bgcolor('black')
        ax.spines['top'].set_edgecolor(frame_color)
        ax.spines['bottom'].set_edgecolor(frame_color)
        ax.spines['left'].set_edgecolor(frame_color)
        ax.spines['right'].set_edgecolor(frame_color)
    #     ax.yaxis.set_major_formatter(kilo_formatter)
    #     ax.xaxis.set_major_formatter(kilo_formatter)
    base_date = datetime.strptime(t.units, "seconds since %Y-%m-%d %H:%M:%S")
    time_delta = timedelta(0,float(t[1]-t[0]),0)
    start_time = base_date + time_delta

    indexer = [None,]*len(grid.shape)

    frame_start_times = []
    for i in range(n_frames):
        frame_start = base_date + timedelta(0,float(t[i]),0)
        frame_start_times.append(frame_start)
        indexer[grid_t_idx] = i

        density = grid[indexer]

        # density,edges = np.histogramdd((x,y), bins=(xedge,yedge))
        density_plot  = p.multiples.flat[i].pcolormesh(xedge,yedge,
                                   np.log10(density.transpose()),
                                   vmin=-0.2,
                                   vmax=np.log10(max_count),
                                   cmap=colormap)
        label_string = frame_start.strftime('%H%M:%S')
        text_label = p.multiples.flat[i].text(xedge[0]-pad+x_range*.015, yedge[0]-pad+y_range*.015, label_string, color=grey_color, size=6)
        density_plot.set_rasterized(True)
        density_maxes.append(density.max())
        total_counts.append(density.sum())
        all_t.append(frame_start)
        print label_string, x.shape, density.max(), density.sum()

    color_scale = ColorbarBase(p.colorbar_ax, cmap=density_plot.cmap,
                                       norm=density_plot.norm,
                                       orientation='horizontal')

    # color_scale.set_label('count per pixel')
    color_scale.set_label('log10(count per pixel)')

    view_x = (xedge.min(), xedge.max())
    view_y = (yedge.min(), yedge.max())

    print 'making multiples',
    p.multiples.flat[0].axis(view_x+view_y)
    filename = '%s-%s_%s_%05.2fkm_%05.1fs.%s' % (filename_prefix, grid_name, start_time.strftime('%Y%m%d_%H%M%S'), dx, time_delta.seconds, image_type)
    filename = os.path.join(outpath, filename)
    if do_save:
        fig.savefig(filename, dpi=150)

    return fig, p, frame_start_times, filename

    print ' ... done'

if __name__ == '__main__':
    do_profile=False
    if do_profile:
        import hotshot
        from hotshot import stats
        prof = hotshot.Profile("multiples_test_profile")
        prof.runcall(runtest)
        prof.close()
        s=stats.load("multiples_test_profile")
        s.sort_stats("time").print_stats()
    else:
        import sys
        res = runtest(sys.argv[1], sys.argv[2])<|endoftext|>


<commit_before># coding: utf-8
import sys, urlparse, csv, cStringIO, codecs

import requests
from lxml import html, etree

PROVINCIAS = {
    'BUE':'Buenos Aires',
    'CAT':'Catamarca',
    'CHA':'Chaco',
    'CHU':'Chubut',
    'CBA':'Córdoba',
    'COR':'Corrientes',
    'ERI':'Entre Ríos',
    'FOR':'Formosa',
    'JUJ':'Jujuy',
    'LAP':'La Pampa',
    'LRJ':'La Rioja',
    'MZA':'Mendoza',
    'MIS':'Misiones',
    'NEU':'Neuquén',
    'RNO':'Río Negro',
    'SAL':'Salta',
    'SJU':'San Juan',
    'SLU':'San Luis',
    'SCR':'Santa Cruz',
    'SFE':'Santa Fe',
    'SGO':'Santiago del Estero',
    'TDF':'Tierra del Fuego',
    'TUC':'Tucumán',
}

FIELDS = [
    'Codigo Municipio',
    'Codigo Provincia',
    'Nombre Municipio',
    'Nombre Provincia',
    'Municipio Lat',
    'Municipio Long',
    u'Datos de Contacto Direcci\xf3n postal',
    'Datos de Contacto E-mail',
    'Datos de Contacto Sitio web',
    u'Datos de Contacto Tel\xe9fonos',
    u'Info Institucional Carta Org\xe1nica',
    u'Info Institucional Categor\xeda',
    u'Info Institucional Fecha de creaci\xf3n',
    'Jefe de Gobierno Cargo',
    'Jefe de Gobierno Nombre y Apellido',
    'Jefe de Gobierno Reelecto'
]

class DictUnicodeWriter(object):

    def __init__(self, f, fieldnames, dialect=csv.excel, encoding="utf-8", **kwds):
        # Redirect output to a queue
        self.fieldnames = fieldnames
        self.queue = cStringIO.StringIO()
        self.writer = csv.DictWriter(self.queue, fieldnames, dialect=dialect, **kwds)
        self.stream = f
        self.encoder = codecs.getincrementalencoder(encoding)()

    def writerow(self, D):
        self.writer.writerow({k:v.encode("utf-8") for k,v in D.items()})
        # Fetch UTF-8 output from the queue ...
        data = self.queue.getvalue()
        data = data.decode("utf-8")
        # ... and reencode it into the target encoding
        data = self.encoder.encode(data)
        # write to the target stream
        self.stream.write(data)
        # empty queue
        self.queue.truncate(0)

    def writerows(self, rows):
        for D in rows:
            self.writerow(D)

    def writeheader(self):
        header = dict(zip(self.fieldnames, self.fieldnames))
        self.writerow(header)

def get_munis_by_prov(prov_code):
    """ genera lista de codigos de municipios segun `prov_code` """
    r = requests.get('http://www.mininterior.gov.ar/municipios/lista_municipios.php?provincia=%s' % prov_code)
    for opt in html.document_fromstring(r.text).cssselect('option'):
        if opt.attrib.get('value', '') != '':
            yield opt.attrib['value'], opt.text

def get_muni_data(code):
    """ scrapea datos de municipio, retorna dict """
    r = requests.get('http://www.mininterior.gov.ar/municipios/masinfo.php?municipio=%s' % code)

    if 'mal formado' in r.text:
        return None

    doc = html.document_fromstring(r.text)

    rv = {}

    rv['Nombre Provincia'] = PROVINCIAS[code[:3]].decode('utf-8')
    rv['Codigo Provincia'] = code[:3]
    rv['Codigo Municipio'] = code
    rv['Nombre Municipio'] = doc.cssselect('div#info h1')[0].text_content().strip()
    rv['Municipio Long'], rv['Municipio Lat'] = '', ''

    iframe = doc.cssselect('iframe')
    if len(iframe) > 0:
        rv['Municipio Long'], rv['Municipio Lat'] = urlparse.parse_qs(urlparse.urlparse(iframe[0].attrib['src']).query)['q'][0].split(',')

    # jefe de gobierno
    table = doc.xpath('//h3/strong[contains(., "Jefe")]/../following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Jefe de Gobierno'))

    # informacion institucional
    table = doc.xpath('//h3[contains(., "Institucional")]/following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Info Institucional'))

    # datos de contacto
    table = doc.xpath('//h3[contains(., "DATOS DE")]/following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Datos de Contacto'))

    return rv

def _scrape_table(element, key_prefix=''):
    """ recibe un `Element`, devuelve diccionario """
    trs = element.cssselect('tr')
    if len(trs) > 1:
        return { key_prefix + ' ' + tr.cssselect('td')[0].text_content().strip(): unicode(tr.cssselect('td')[1].text_content().strip())
                 for tr in element.cssselect('tr') }
    else:
        return {}

def main():
    csv_writer = DictUnicodeWriter(sys.stdout, FIELDS)
    csv_writer.writeheader()
    for prov_code in PROVINCIAS.keys():
        for muni_code, _ in get_munis_by_prov(prov_code):
            print >>sys.stderr, '---'
            print >>sys.stderr, "PARSING MUNI: %s" % muni_code
            row = get_muni_data(muni_code)
            if row is None:
                print >>sys.stderr, 'Error en fuente'
                print >>sys.stderr,'---'
                continue
            print >>sys.stderr,row
            csv_writer.writerow(row)
            print >>sys.stderr,'---'

if __name__ == '__main__':
    main()
<commit_msg>Arreglar selector css<commit_after># coding: utf-8
import sys, urlparse, csv, cStringIO, codecs

import requests
from lxml import html, etree

PROVINCIAS = {
    'BUE':'Buenos Aires',
    'CAT':'Catamarca',
    'CHA':'Chaco',
    'CHU':'Chubut',
    'CBA':'Córdoba',
    'COR':'Corrientes',
    'ERI':'Entre Ríos',
    'FOR':'Formosa',
    'JUJ':'Jujuy',
    'LAP':'La Pampa',
    'LRJ':'La Rioja',
    'MZA':'Mendoza',
    'MIS':'Misiones',
    'NEU':'Neuquén',
    'RNO':'Río Negro',
    'SAL':'Salta',
    'SJU':'San Juan',
    'SLU':'San Luis',
    'SCR':'Santa Cruz',
    'SFE':'Santa Fe',
    'SGO':'Santiago del Estero',
    'TDF':'Tierra del Fuego',
    'TUC':'Tucumán',
}

FIELDS = [
    'Codigo Municipio',
    'Codigo Provincia',
    'Nombre Municipio',
    'Nombre Provincia',
    'Municipio Lat',
    'Municipio Long',
    u'Datos de Contacto Direcci\xf3n postal',
    'Datos de Contacto E-mail',
    'Datos de Contacto Sitio web',
    u'Datos de Contacto Tel\xe9fonos',
    u'Info Institucional Carta Org\xe1nica',
    u'Info Institucional Categor\xeda',
    u'Info Institucional Fecha de creaci\xf3n',
    'Jefe de Gobierno Cargo',
    'Jefe de Gobierno Nombre y Apellido',
    'Jefe de Gobierno Reelecto'
]

class DictUnicodeWriter(object):

    def __init__(self, f, fieldnames, dialect=csv.excel, encoding="utf-8", **kwds):
        # Redirect output to a queue
        self.fieldnames = fieldnames
        self.queue = cStringIO.StringIO()
        self.writer = csv.DictWriter(self.queue, fieldnames, dialect=dialect, **kwds)
        self.stream = f
        self.encoder = codecs.getincrementalencoder(encoding)()

    def writerow(self, D):
        self.writer.writerow({k:v.encode("utf-8") for k,v in D.items()})
        # Fetch UTF-8 output from the queue ...
        data = self.queue.getvalue()
        data = data.decode("utf-8")
        # ... and reencode it into the target encoding
        data = self.encoder.encode(data)
        # write to the target stream
        self.stream.write(data)
        # empty queue
        self.queue.truncate(0)

    def writerows(self, rows):
        for D in rows:
            self.writerow(D)

    def writeheader(self):
        header = dict(zip(self.fieldnames, self.fieldnames))
        self.writerow(header)

def get_munis_by_prov(prov_code):
    """ genera lista de codigos de municipios segun `prov_code` """
    r = requests.get('http://www.mininterior.gov.ar/municipios/lista_municipios.php?provincia=%s' % prov_code)
    for opt in html.document_fromstring(r.text).cssselect('option'):
        if opt.attrib.get('value', '') != '':
            yield opt.attrib['value'], opt.text

def get_muni_data(code):
    """ scrapea datos de municipio, retorna dict """
    r = requests.get('http://www.mininterior.gov.ar/municipios/masinfo.php?municipio=%s' % code)

    if 'mal formado' in r.text:
        return None

    doc = html.document_fromstring(r.text)

    rv = {}

    rv['Nombre Provincia'] = PROVINCIAS[code[:3]].decode('utf-8')
    rv['Codigo Provincia'] = code[:3]
    rv['Codigo Municipio'] = code
    rv['Nombre Municipio'] = doc.cssselect('div#mas-info h1')[0].text_content().strip()
    rv['Municipio Long'], rv['Municipio Lat'] = '', ''

    iframe = doc.cssselect('iframe')
    if len(iframe) > 0:
        rv['Municipio Long'], rv['Municipio Lat'] = urlparse.parse_qs(urlparse.urlparse(iframe[0].attrib['src']).query)['q'][0].split(',')

    # jefe de gobierno
    table = doc.xpath('//h3/strong[contains(., "Jefe")]/../following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Jefe de Gobierno'))

    # informacion institucional
    table = doc.xpath('//h3[contains(., "Institucional")]/following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Info Institucional'))

    # datos de contacto
    table = doc.xpath('//h3[contains(., "DATOS DE")]/following-sibling::table[1]')[0]
    rv.update(_scrape_table(table, 'Datos de Contacto'))

    return rv

def _scrape_table(element, key_prefix=''):
    """ recibe un `Element`, devuelve diccionario """
    trs = element.cssselect('tr')
    if len(trs) > 1:
        return { key_prefix + ' ' + tr.cssselect('td')[0].text_content().strip(): unicode(tr.cssselect('td')[1].text_content().strip())
                 for tr in element.cssselect('tr') }
    else:
        return {}

def main():
    csv_writer = DictUnicodeWriter(sys.stdout, FIELDS)
    csv_writer.writeheader()
    for prov_code in PROVINCIAS.keys():
        for muni_code, _ in get_munis_by_prov(prov_code):
            print >>sys.stderr, '---'
            print >>sys.stderr, "PARSING MUNI: %s" % muni_code
            row = get_muni_data(muni_code)
            if row is None:
                print >>sys.stderr, 'Error en fuente'
                print >>sys.stderr,'---'
                continue
            print >>sys.stderr,row
            csv_writer.writerow(row)
            print >>sys.stderr,'---'

if __name__ == '__main__':
    main()
<|endoftext|>