<commit_before>
# This Python file uses the following encoding: utf-8
# Recovery docs have their own system for ids.
# Annual reports before 2006 are not included because of broken links and unique formatting
# Still need a plan for "indexed" forms
import re
from bs4 import BeautifulSoup
from datetime import datetime, date
from utils import utils, inspector

report = {}
source_links = {}
agency_decoder = {
    "Department of Justice":["Department of Justice", "doj"],
    "United States Marshals Service (USMS)": [ "United States Marshals Service", "USMS"],
    "Office of Justice Programs (OJP)": ["Office of Justice Programs", "OJP"],
    "Federal Bureau of Prisons (BOP)": ["Federal Bureau of Prisons", "BOP"],
    "Federal Bureau of Investigation (FBI)": ["Federal Bureau of Investigation", "FBI"],
    "Combined DNA Index System (CODIS)": ["Combined DNA Index System", "CODIS"],
    "Executive Office for U. S. Attorneys (EOUSA)": ["Executive Office for U. S. Attorneys", "EOUSA"],
    "Office of Community Oriented Policing Services (COPS)": ["Office of Community Oriented Policing Services", "COPS"],
    "Intergovernmental Agreements for Detention Space (IGAs)": ["Intergovernmental Agreements for Detention Space", "IGAs"],
    "Drug Enforcement Administration (DEA)": ["Drug Enforcement Administration", "DEA"],
    "Bureau of Alcohol, Tobacco, Firearms and Explosives (ATF)": ["Bureau of Alcohol, Tobacco, Firearms and Explosives", "ATF"],
    "Office on Violence Against Women (OVW)": ["Office on Violence Against Women", "OVW"],
    "Immigration and Naturalization Service (INS) – 1994 to 2003": ["Immigration and Naturalization Service", "INS"],
    "United States Marshals Service (USMS)": ["United States Marshals Service", "USMS"],
  }

base_url = "http://www.justice.gov"
not_agency = ("Office of Justice Programs (OJP)", "Contracts", "Special Reports", "Other DOJ Components and Reports Encompassing More Than One DOJ Component","Equitable Sharing", "Offices, Boards and Divisions (OBDs)")
index_links = ["http://www.justice.gov/oig/reports/2012/s1209.htm"]

def extract_info(content, directory):
  if directory in not_agency:
    agency = "doj"
    agency_name = "Department of Justice"
  elif directory[:11] == "Immigration":
    agency = "ins"
    agency_name = "Immigration and Naturalization Service"
  else:
    agency = agency_decoder[directory][1]
    agency_name = agency_decoder[directory][0]

  for c in content:
    blurbs = c.find_all("p")

  for b in blurbs:
    # date
    x = b.previous_sibling
    y = b.previous_sibling.previous_sibling
    try:
      if y['class'] == ['date']:
        date_string = y.string
      elif x['class'] == ['date']:
        date_string = x.string
    except:
      date_string = None

    if date_string == None:
      try:
        date = b.string
        # get rid of the last (...)
        date = re.sub(r'\([^)]*\)', '', date)
        # get rid of the last [...]
        date = re.sub(r'\[[^)]*\]', '', date)
        # last chunk after the ","
        date = date.rsplit(',')
        date = date[-1]
        date_string = date

        # date = datetime.strptime(date, "YYYY-MM-DD")
        # published_on = date
      except:
        date_string = None

    if date_string == None:
      date = b.contents[0].string
      date = date.replace(" [", "")
      date = date.replace("[", "")
      date = date.rsplit(',')
      date = date[-1]
      date_string = date

    date_string = date_string.strip()

    if "," not in date_string:
      date_string = date_string.replace(" ", " 1, ")

    date = datetime.strptime(date_string, "%B %d, %Y")
    year = datetime.strftime(date, "%Y")
    published_on = datetime.strftime(date, "%Y-%m-%d")

    string_title = b.string
    if b.string == None:
      string_title = b.contents
      if "<a href=" in str(string_title):
        string_title = b.contents[0]

    for l in b.find_all("a"):
      link = l.get("href")
      if link != None:
        # title
        title = l.string
        if title == "HTML" or title == "PDF":
          title = string_title
        if title == None:
          #this is the only one that doesn't work
          if str(b) == '
Report on Activities Under Section 702 of the FISA Amendments Act of 2008
':
            title = "Report on Activities Under Section 702 of the FISA Amendments Act of 2008"

        if "[" in title[-5:]:
          title = title.replace("[", '')
        title = title.strip()

        # formating links consistently
        if link[:1] != "/":
          link = "/" + link
        # id
        doc_id = str(link)[1:-4]
        if link[:11] == "oig/reports":
          doc_id = doc_id[16:]

        #these docs are one report where the page has a table of contents with links to content
        if "/index" in link:
          index_links.append(link)
          indexed = True
        else:
          indexed = False

        # there may be a better way to do this but I am just taking out all the things that are not the id
        url_extras = ( "/final", "/press", "/fullpdf", "/ins_response", "oig/special/", "USMS/", "plus/", "oig/grants/", "oig/reports/", "EOUSA/", "BOP/", "ATF/", "COPS/", "FBI/", "OJP/", "INS/", "DEA/", "OBD", "Final/", "full/", "/analysis", "/report", "/PDF_list", "/exec", "/full_report", "/full", "_redacted", "oig", "r-summary", "/response", "/listpdf", "/memo", "/fullreport", "/Final", "/extradition")
        for n in url_extras:
          if n in doc_id:
            doc_id = doc_id.replace(n, "")

        if doc_id[:1] == "/":
          doc_id = doc_id[1:]

        if "/" in doc_id:
          if doc_id[4:5] == "/":
            if doc_id[:2] == "19" or doc_id[:2] == "20":
              doc_id = doc_id[5:]

        special_cases = {"a0118/au0118":"a0118", "a0207/0207":"a0207",  }
        if doc_id in special_cases.keys():
          doc_id = special_cases[doc_id]

        # url
        if link[:5] == "/oig/":
          url = base_url + link
        else:
          url = base_url + "/oig/reports" + link
        # file_type
        if url[-3:] == "pdf":
          file_type = "pdf"
        elif url[-3:] == "htm":
          file_type = "html"
        elif url[-4:] == "html":
          file_type = "html"
        else:
          # these include a few navigation links
          file_type = "ignore"

        if file_type != "ignore":
          if report.has_key(doc_id):
            if file_type == "pdf":
              # current and previous file pdf
              if report[doc_id]["file_type"] == "pdf":
                report[doc_id]["categories"].append(directory)
              # current file a pdf, old file html
              else:
                report[doc_id]["file_type"] = "pdf"
                report[doc_id]["url"] = url
                report[doc_id]["categories"].append(directory)
            else:
              # current file html old file pdf OR both files html
              report[doc_id]["categories"].append(directory)

            # add url if new
            for n in report[doc_id]["urls"]:
              if n.has_key(url):
                old_url = True
            if not "old_url" in locals():
              report[doc_id]["urls"].append({
                "url":url,
                "file_type": file_type,
                "indexed": indexed,
                })

            # finding the most descriptive name for cross-listed docs
            if report[doc_id]["agency"] == "doj" and agency != "doj":
              report[doc_id]["agency"] = agency
              report[doc_id]["agency_name"] = agency_name

          else:
            report[doc_id] = {
              "report_id": doc_id,
              "inspector": "doj",
              "inspector_url": "http://www.justice.gov/oig/reports/",
              "agency": agency,
              "agency_name": agency_name,
              "url": url,
              "title": title,
              "file_type": file_type,
              "categories": [directory,],
              "urls": [{
                  "url":url,
                  "file_type": file_type,
                  "indexed": indexed,
                }],
              "published_on": published_on,
              "year": year,
              # perhaps elaborate on this later
              "type": "report",
              }

def get_content(url):
  page = utils.download(url)
  page = BeautifulSoup(page)
  content = page.select(".content-left")
  return content

def find_pages():
  starting_point = "http://www.justice.gov/oig/reports/"
  content = get_content(starting_point)
  for c in content:
    links = c.find_all("a")
    for l in links:
      name = l.string
      link = base_url + l.get("href")
      source_links[link] = name

def run():
  find_pages()
  for l in source_links.keys():
    content = get_content(l)
    extract_info(content, source_links[l])

  # # for debugging I am using a test file
  # f = open("inspectors/USDOJscraper_practice.html","r")
  # data = f.read()
  # page = BeautifulSoup(data)
  # content = page.select(".content-left")
  # extract_info(content, "Bureau of Alcohol, Tobacco, Firearms and Explosives (ATF)")

  for key in report.keys():
    inspector.save_report(report[key])

run()

<commit_msg>I almost have it working locally but still catching edge cases here and there<commit_after>
# This Python file uses the following encoding: utf-8
# Recovery docs have their own system for ids.
# Annual reports before 2006 are not included because of broken links and unique formatting
# Still need a plan for "indexed" forms
import re
from bs4 import BeautifulSoup
from datetime import datetime, date
from utils import utils, inspector

report = {}
source_links = {}
agency_decoder = {
    "Department of Justice":["Department of Justice", "doj"],
    "United States Marshals Service (USMS)": [ "United States Marshals Service", "USMS"],
    "Office of Justice Programs (OJP)": ["Office of Justice Programs", "OJP"],
    "Federal Bureau of Prisons (BOP)": ["Federal Bureau of Prisons", "BOP"],
    "Federal Bureau of Investigation (FBI)": ["Federal Bureau of Investigation", "FBI"],
    "Combined DNA Index System (CODIS)": ["Combined DNA Index System", "CODIS"],
    "Executive Office for U. S. Attorneys (EOUSA)": ["Executive Office for U. S. Attorneys", "EOUSA"],
    "Office of Community Oriented Policing Services (COPS)": ["Office of Community Oriented Policing Services", "COPS"],
    "Intergovernmental Agreements for Detention Space (IGAs)": ["Intergovernmental Agreements for Detention Space", "IGAs"],
    "Drug Enforcement Administration (DEA)": ["Drug Enforcement Administration", "DEA"],
    "Bureau of Alcohol, Tobacco, Firearms and Explosives (ATF)": ["Bureau of Alcohol, Tobacco, Firearms and Explosives", "ATF"],
    "Office on Violence Against Women (OVW)": ["Office on Violence Against Women", "OVW"],
    "Immigration and Naturalization Service (INS) – 1994 to 2003": ["Immigration and Naturalization Service", "INS"],
    "United States Marshals Service (USMS)": ["United States Marshals Service", "USMS"],
  }

base_url = "http://www.justice.gov"
not_agency = ("Office of Justice Programs (OJP)", "Contracts", "Special Reports", "Other DOJ Components and Reports Encompassing More Than One DOJ Component","Equitable Sharing", "Offices, Boards and Divisions (OBDs)")
index_links = ["http://www.justice.gov/oig/reports/2012/s1209.htm"]

def extract_info(content, directory):
  if directory in not_agency:
    agency = "doj"
    agency_name = "Department of Justice"
  elif directory[:11] == "Immigration":
    agency = "ins"
    agency_name = "Immigration and Naturalization Service"
  else:
    agency = agency_decoder[directory][1]
    agency_name = agency_decoder[directory][0]

  for c in content:
    blurbs = c.find_all("p")

  for b in blurbs:
    # date
    x = b.previous_sibling
    y = b.previous_sibling.previous_sibling
    try:
      if y['class'] == ['date']:
        date_string = y.string
      elif x['class'] == ['date']:
        date_string = x.string

    except:
      date_string = None

    if date_string == None:
      try:
        date = b.string
        date = re.sub(r'\([^)]*\)', '', date)
        date = re.sub(r'\[(.*?)\]', '', date)
        date = date.rsplit(',')
        print date, 1
        date = date[-1]
        print date, 2
        date_string = date.strip()
        print date , 3

      except:
        date_string = None

    if date_string == None:
      date = b.contents[0].string
      date = re.sub(r'\([^)]*\)', '', date)
      date = re.sub(r'\[(.*?)\]', '', date)
      date = date.replace(" [", "")
      date = date.replace("[", "")
      date = date.rsplit(',')
      date = date[-1]
      date_string = date

    print date_string, "THIS ONE"
    date_string = date_string.strip()

    if "," not in date_string:
      date_string = date_string.replace(" ", " 1, ")

    # going through each link in a paragraph
    for l in b.find_all("a"):
      try:
        date = datetime.strptime(date_string, "%B %d, %Y")
      except:
        print l
        print date
        info = odd_link(b, date)
        real_title = info["real_title"]
        date_string = info["date_string"]
        date = datetime.strptime(date_string, "%B %d, %Y")

      year = datetime.strftime(date, "%Y")
      published_on = datetime.strftime(date, "%Y-%m-%d")

      string_title = b.string
      if b.string == None:
        string_title = b.contents
        if "<a href=" in str(string_title):
          string_title = b.contents[0]

      #for l in b.find_all("a"):
      link = l.get("href")
      if link != None:
        # title
        title = l.string
        if title == "HTML" or title == "PDF":
          title = string_title
        if title == None:
          #this is the only one that doesn't work
          if str(b) == '
Report on Activities Under Section 702 of the FISA Amendments Act of 2008
':
            title = "Report on Activities Under Section 702 of the FISA Amendments Act of 2008"

        if "real_title" in locals():
          title = real_title

        if "[" in title[-5:]:
          title = title.replace("[", '')
        title = title.strip()

        # formating links consistently
        if link[:1] != "/":
          link = "/" + link
        # id
        doc_id = str(link)[1:-4]
        if link[:11] == "oig/reports":
          doc_id = doc_id[16:]

        #these docs are one report where the page has a table of contents with links to content
        if "/index" in link:
          index_links.append(link)
          indexed = True
        else:
          indexed = False

        # there may be a better way to do this but I am just taking out all the things that are not the id
        url_extras = ( "/final", "/fullpdf", "/ins_response", "oig/special/", "USMS/", "plus/", "oig/grants/", "oig/reports/", "EOUSA/", "BOP/", "ATF/", "COPS/", "FBI/", "OJP/", "INS/", "DEA/", "OBD", "/analysis", "/report", "/PDF_list", "/full_report", "/full", "_redacted", "oig", "r-", "/response", "/listpdf", "/memo", "/fullreport", "/Final", "/extradition")
        for n in url_extras:
          if n in doc_id:
            doc_id = doc_id.replace(n, "")

        if doc_id[:1] == "/":
          doc_id = doc_id[1:]

        if "/" in doc_id:
          if doc_id[4:5] == "/":
            if doc_id[:2] == "19" or doc_id[:2] == "20":
              doc_id = doc_id[5:]

        special_cases = {"a0118/au0118":"a0118", "a0207/0207":"a0207",  }
        if doc_id in special_cases.keys():
          doc_id = special_cases[doc_id]

        if "spanish" in link:
          language = "Spanish"
        else:
          language = "English"

        # url
        if link[:5] == "/oig/":
          url = base_url + link
        else:
          url = base_url + "/oig/reports" + link
        # file_type
        file_type = find_file_type(link)

        if file_type != "ignore":
          if report.has_key(doc_id):
            if file_type == "pdf":
              # current and previous file pdf
              if report[doc_id]["file_type"] == "pdf":
                report[doc_id]["categories"].append(directory)
              # current file a pdf, old file html
              else:
                report[doc_id]["file_type"] = "pdf"
                report[doc_id]["url"] = url
                report[doc_id]["categories"].append(directory)
            else:
              # current file html old file pdf OR both files html
              report[doc_id]["categories"].append(directory)

            # add url if new
            for n in report[doc_id]["urls"]:
              if n.has_key(url):
                old_url = True
            if not "old_url" in locals():
              report[doc_id]["urls"].append({
                "url":url,
                "file_type": file_type,
                "indexed": indexed,
                })

            # finding the most descriptive name for cross-listed docs
            if report[doc_id]["agency"] == "doj" and agency != "doj":
              report[doc_id]["agency"] = agency
              report[doc_id]["agency_name"] = agency_name

          else:
            report[doc_id] = {
              "report_id": doc_id,
              "inspector": "doj",
              "inspector_url": "http://www.justice.gov/oig/reports/",
              "agency": agency,
              "agency_name": agency_name,
              "url": url,
              "title": title,
              "file_type": file_type,
              "categories": [directory,],
              "urls": [{
                  "url":url,
                  "file_type": file_type,
                  "indexed": indexed,
                }],
              "published_on": published_on,
              "year": year,
              # perhaps elaborate on this later
              "type": "report",
              "language": language,
              }

def find_file_type(url):
  if url[-3:] == "pdf":
    file_type = "pdf"
  elif url[-3:] == "htm":
    file_type = "html"
  elif url[-4:] == "html":
    file_type = "html"
  else:
    # these include a few navigation links
    file_type = "ignore"
  return file_type

def date_format(date):
  date = str(date)
  date = re.sub(r'\([^)]*\)', '', date)
  date = re.sub(r'\[(.*?)\]', '', date)
  date = date.rsplit(',')
  date = str(date[-2]).strip() + ", " + str(date[-1]).strip()
  date_string = date
  return date_string

def odd_link(b, date):
  #print date, "hello", b
  if date.strip() =="Report of Investigation into Allegations Relating to the Selection of the":
    return {"date_string":"June 1, 2006", "real_title": "Report of Investigation into Allegations Relating to the Selection of the U. S. Attorney for Guam and the Northern Mariana Islands, June 2006" }
  if "Released Publicly" in date:
    date =  date[18:].strip()
    date = date.replace(" ", " 1, ")
    t = str(b)
    t = t[4:-4]
    return{"date_string": date, "real_title": t}
  if "(Unclassified Summary)" in date:
    date = str(b)
    date = date[4:-4]
    date = re.sub(r'\([^)]*\)', '', date)
    date = re.sub(r'\[(.*?)\]', '', date)
    date = date.rsplit(',')
    date = date[-1]
    if "\r\n " in date:
      date = re.sub(r'\r\n ', '', date)
      date = date.strip()
      date = date.replace(" ", " 1, ")
    t = str(b)
    t = t[4:-4]
    t = re.sub(r'\[(.*?)\]', '', t)
    return{"date_string": date, "real_title": t}
  if "" in str(b):
    print "hello"
    print date
    text = str(b)[3:-4]
    text = text.replace("", "")
    text = text.replace("", "")
    title = text.replace("
", "")
    title = title.replace("\r\n", "")
    date = re.sub(r'\([^)]*\)', '', title)
    date = re.sub(r'\[[^)]*\]', '', date)
    date = date.rsplit(',')
    #print date
    date_string = date[-1]
    date_string = date_string.strip()
    #print date_string
    if "," not in date_string:
      date_string = date_string.replace(" ", " 1, ")

    return{"date_string": date_string, "real_title": title}
  if date.strip() == "Alleged Deception of Congress: The Congressional Task Force on Immigration Reform's Fact-Finding Visit to the Miami District of INS in June 1995":
    print "WORKING"
    return{"date_string": "June 1, 1996", "real_title": "Alleged Deception of Congress: The Congressional Task Force on Immigration Reform's Fact-Finding Visit to the Miami District of INS in June 1995"}
  if date != None:
    date = date.strip
    # case 1, date is wrong because it is in the paragraph and completely written out
    try:
        date =  b.string
        date_string = date_format(date)
        title = b.string
    except:
      # these are lists of links that are different variants of the same report in a list
      print "here"
      # case where there is a list in a paragraph tag
      listy = b.parent.parent
      text = str(listy.previous_sibling)
      if "


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

from httplib2 import Http

import os
import urllib
import urllib2
import base64
import json
import sys
import argparse
import time

try:
    import requests
except ImportError:
    print 'The requests package is required: http://docs.python-requests.org/en/latest/user/install/#install'
    sys.exit()

try:
    import json
except ImportError:
    import simplejson as json

SQL_URL = "%(protocol)s://%(user)s.%(domain)s/api/%(sql_version)s/sql"
IMPORT_URL = '%(protocol)s://%(user)s.%(domain)s/api/%(import_version)s/imports/%(queue_id)s?api_key=%(api_key)s'

class CartoDBDashboardException(Exception):
    pass

class CartoDbDashboard:
    def __init__(self, user, key, host='cartodb.com', protocol='https', sql_version='v2', import_version='v1', verbose=True):

        self.user = user
        self.key = key
        self.host = host
        self.protocol = protocol
        self.sql_version = sql_version
        self.import_version = import_version
        self.verbose = verbose

        self.sql_url = SQL_URL % {'user': user, 'domain': host, 'protocol': protocol, 'sql_version': sql_version}
        self.import_url = IMPORT_URL % {'user': user, 'domain': host, 'protocol': protocol,'import_version': import_version, 'api_key': key, 'queue_id':''}

    def _log(self, message):
        if self.verbose == True:
            print message

    def _error(self, error):
        print error
        sys.exit()

    def sql_api(self, sql):
        # execute sql request over API
        try:
            params = {
                'api_key' : self.key,
                'q'       : sql
            }
            r = requests.get(self.sql_url, params=params)
            return r.json()
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))

    def convert_data_type(self, column, coltype, table):
        try:
            if coltype == 'string':
                coltype = 'text'
            elif coltype == 'number':
                coltype = 'numeric USING NULLIF( %s , \'\')::numeric' %(column)

            sql = "ALTER TABLE %s ALTER COLUMN %s TYPE %s" % (table, column, coltype)
            data = self.sql_api(sql)
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def import_data(self, datafile, type_guessing = 'true'):
        try:
            self._log('NEW..')
            params = {
                'type_guessing' : type_guessing
            }

            r = requests.post(self.import_url, files={'file': open(datafile, 'rb')}, params=params)

            data = r.json()
            if data['success']!=True:
                self._error("Upload failed")
            complete = False
            last_state = ''
            while not complete:
                import_status_url = IMPORT_URL % {'user': self.user, 'domain': self.host, 'protocol': self.protocol, 'sql_version': self.sql_version, 'import_version': self.import_version, 'api_key': self.key,'queue_id':data['item_queue_id']}
                req = urllib2.Request(import_status_url)
                response = urllib2.urlopen(req)
                d = json.loads(str(response.read()))
                if last_state!=d['state']:
                    last_state=d['state']
                    if d['state']=='uploading':
                        self._log('Uploading file...')
                    elif d['state']=='importing':
                        self._log('Importing data...')
                    elif d['state']=='complete':
                        if d['success']==True:
                            complete = True
                            self._log('Table "%s" created' % d['table_name'])
                            return complete,d['table_name']
                        else:
                            self._error(d['get_error_text']['what_about'])

                if d['state']=='failure':
                    self._error(d['get_error_text']['what_about'])

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def drop_table(self, table):
        # drop a table '
        try:
            self._log("Dropping table %s"  % table)
            sql = "DROP TABLE %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def clear_rows(self, table):
        # clear all rows from a table
        try:
            self._log("Deleting all rows")
            sql = "DELETE FROM %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def clean_table(self,table):
        # clean up table for speed
        try:
            self._log("Cleaning up unused space")
            sql = "VACUUM FULL %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            self._log("Optimizing existing indexes")
            sql = "ANALYZE %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def get_row_count(self,table):
        try:
            count = -1
            sql = "SELECT count(*) FROM %s" % table
            data = self.sql_api(sql)
            if data:
                if 'error' in data.keys():
                    return count
                count = data['rows'][0]['count']
            return count
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return 0

    def rename_table(self, table_name, new_table_name):
        print "Start renaming process by a little wait of 30s from %s" % time.ctime()
        time.sleep( 60 )
        print "Hopefully we are now good to go"

        try:
            if self.table_exists(table_name):
                sql = "ALTER TABLE %s  RENAME TO %s" % (table_name, new_table_name)
                data = self.sql_api(sql)
                if 'error' in data.keys():
                    self._error(data['error'])
                return True
            return False

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def table_exists(self, table_name):
        try:
            if self.get_row_count(table_name) >=0:
                return True
            return False

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False
<commit_msg>change carto to host name and bump version to 0.2.12<commit_after># -*- encoding: utf-8 -*-

from httplib2 import Http

import os
import urllib
import urllib2
import base64
import json
import sys
import argparse
import time

try:
    import requests
except ImportError:
    print 'The requests package is required: http://docs.python-requests.org/en/latest/user/install/#install'
    sys.exit()

try:
    import json
except ImportError:
    import simplejson as json

SQL_URL = "%(protocol)s://%(user)s.%(domain)s/api/%(sql_version)s/sql"
IMPORT_URL = '%(protocol)s://%(user)s.%(domain)s/api/%(import_version)s/imports/%(queue_id)s?api_key=%(api_key)s'

class CartoDBDashboardException(Exception):
    pass

class CartoDbDashboard:
    def __init__(self, user, key, host='carto.com', protocol='https', sql_version='v2', import_version='v1', verbose=True):

        self.user = user
        self.key = key
        self.host = host
        self.protocol = protocol
        self.sql_version = sql_version
        self.import_version = import_version
        self.verbose = verbose

        self.sql_url = SQL_URL % {'user': user, 'domain': host, 'protocol': protocol, 'sql_version': sql_version}
        self.import_url = IMPORT_URL % {'user': user, 'domain': host, 'protocol': protocol,'import_version': import_version, 'api_key': key, 'queue_id':''}

    def _log(self, message):
        if self.verbose == True:
            print message

    def _error(self, error):
        print error
        sys.exit()

    def sql_api(self, sql):
        # execute sql request over API
        try:
            params = {
                'api_key' : self.key,
                'q'       : sql
            }
            r = requests.get(self.sql_url, params=params)
            return r.json()
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))

    def convert_data_type(self, column, coltype, table):
        try:
            if coltype == 'string':
                coltype = 'text'
            elif coltype == 'number':
                coltype = 'numeric USING NULLIF( %s , \'\')::numeric' %(column)

            sql = "ALTER TABLE %s ALTER COLUMN %s TYPE %s" % (table, column, coltype)
            data = self.sql_api(sql)
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def import_data(self, datafile, type_guessing = 'true'):
        try:
            self._log('NEW..')
            params = {
                'type_guessing' : type_guessing
            }

            r = requests.post(self.import_url, files={'file': open(datafile, 'rb')}, params=params)

            data = r.json()
            if data['success']!=True:
                self._error("Upload failed")
            complete = False
            last_state = ''
            while not complete:
                import_status_url = IMPORT_URL % {'user': self.user, 'domain': self.host, 'protocol': self.protocol, 'sql_version': self.sql_version, 'import_version': self.import_version, 'api_key': self.key,'queue_id':data['item_queue_id']}
                req = urllib2.Request(import_status_url)
                response = urllib2.urlopen(req)
                d = json.loads(str(response.read()))
                if last_state!=d['state']:
                    last_state=d['state']
                    if d['state']=='uploading':
                        self._log('Uploading file...')
                    elif d['state']=='importing':
                        self._log('Importing data...')
                    elif d['state']=='complete':
                        if d['success']==True:
                            complete = True
                            self._log('Table "%s" created' % d['table_name'])
                            return complete,d['table_name']
                        else:
                            self._error(d['get_error_text']['what_about'])

                if d['state']=='failure':
                    self._error(d['get_error_text']['what_about'])

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def drop_table(self, table):
        # drop a table '
        try:
            self._log("Dropping table %s"  % table)
            sql = "DROP TABLE %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def clear_rows(self, table):
        # clear all rows from a table
        try:
            self._log("Deleting all rows")
            sql = "DELETE FROM %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def clean_table(self,table):
        # clean up table for speed
        try:
            self._log("Cleaning up unused space")
            sql = "VACUUM FULL %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            self._log("Optimizing existing indexes")
            sql = "ANALYZE %s" % table
            data = self.sql_api(sql)
            if 'error' in data.keys():
                self._error(data['error'])
            return True
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def get_row_count(self,table):
        try:
            count = -1
            sql = "SELECT count(*) FROM %s" % table
            data = self.sql_api(sql)
            if data:
                if 'error' in data.keys():
                    return count
                count = data['rows'][0]['count']
            return count
        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return 0

    def rename_table(self, table_name, new_table_name):
        print "Start renaming process by a little wait of 30s from %s" % time.ctime()
        time.sleep( 60 )
        print "Hopefully we are now good to go"

        try:
            if self.table_exists(table_name):
                sql = "ALTER TABLE %s  RENAME TO %s" % (table_name, new_table_name)
                data = self.sql_api(sql)
                if 'error' in data.keys():
                    self._error(data['error'])
                return True
            return False

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False

    def table_exists(self, table_name):
        try:
            if self.get_row_count(table_name) >=0:
                return True
            return False

        except CartoDBDashboardException as e:
            self._log('some error occurred: %s' %(e))
            return False
<|endoftext|>


<commit_before>#!/usr/bin/env python

import datetime
import logging
import os
from urllib.parse import urljoin
import re

from bs4 import BeautifulSoup
from utils import utils, inspector

# http://www.loc.gov/about/office-of-the-inspector-general/
# Oldest report: 2001

# options:
#   standard since/year options for a year range to fetch from.
#
# Notes for IG's web team:
#

REPORTS_BY_YEAR_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/annual-reports/'
REPORTS_TESTIMONY_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/congressional-testimony/'
REPORTS_PEER_REVIEWS_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/peer-reviews/'

OTHER_REPORTS = {
  'http://lcweb2.loc.gov/master/libn/about/office-of-the-inspector-general/documents/LCR211-6-oct2010.pdf': '2010-10-01'
}

REPORT_ID_TO_EST_DATE = {
  'TLHAOTBMPIAMECT': datetime.datetime(2013, 3, 1),
  'WTTSOTGPARMA': datetime.datetime(2013, 3, 1),
  'MDCITSOICBISCFPR': datetime.datetime(2010, 10, 1),
  'RRSASBRJ2T2A': datetime.datetime(2007, 7, 1),
}

PARTS_MONTH = ('(January|February|March|April|May|June|July|August|September|'
               'October|November|December)')
RE_PUBLIC_RELEASE = re.compile('Not for Public Release', re.I)
RE_DATE = re.compile(r'%s +\d\d\d\d' % PARTS_MONTH)
RE_SHORT_DATE = re.compile(r'(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sept|Oct|Nov|Dec)'
                           r' +\d\d\d\d')
RE_LONG_DATE = re.compile(r'%s \d\d,? +\d\d\d\d' % PARTS_MONTH)
RE_FISCAL_YEAR = re.compile('Fiscal Year (\d\d\d\d)')
RE_OWLC = re.compile('Open +World +Leadership +Center')
RE_JMCF = re.compile('James +Madison +Council +Fund')

def generate_id(title, published_on=None):
  """Generates a report id given the report title and published date.

  This generally tries to be clever and uses the first letter of each name of
  the title, then the month date and year of the report as the ID.
  """
  words = title.split()
  letters = [word[0] for word in words if re.match(r'\w', word[0])]
  slug = ''.join(letters).upper()
  if published_on:
    return '%s-%s' % (slug[:16], published_on.strftime('%Y-%m-%d'))
  else:
    return slug

class LibraryOfCongressScraper(object):
  def __init__(self, options):
    self.options = options
    self.year_range = inspector.year_range(options)

  def run(self):
    self.get_reports_by_year()

    # Get congressional testimony.
    self.get_listed_reports(REPORTS_TESTIMONY_URL)

    # Get peer reviews.
    self.get_listed_reports(REPORTS_PEER_REVIEWS_URL)

  def get_reports_by_year(self):
    # This page contains semianual reports as well as a few Audit reports for
    # Fiscal Year 2014 and links to sub-pages that contain links for other
    # fiscal years.
    doc = BeautifulSoup(utils.download(REPORTS_BY_YEAR_URL))

    # Get the semiannual reports to Congress.
    self.get_semiannual_reports_to_congress(doc)

    # Reports that are 'bare' on the page, listed explicitly
    bare_report_ul, link_subpage_ul = self.get_uls_past_audit_header(doc)
    self.get_bare_reports(bare_report_ul)

    # Links on the page to audit reports from past fiscal years
    for li in link_subpage_ul.find_all('li'):
      link = li.find('a')
      if link:
        next_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
        doc = BeautifulSoup(utils.download(next_url))
        uls = self.get_uls_past_audit_header(doc)
        assert len(uls) == 1, ('Mysterious additional ul data on page: %s' %
                               next_url)
        self.get_bare_reports(uls[0])

  def get_listed_reports(self, url):
    doc = BeautifulSoup(utils.download(url))
    article = doc.select('.article')[0]
    for ul in article.find_all('ul'):
      self.get_bare_reports(ul)

  def get_semiannual_reports_to_congress(self, doc):
    header = doc.find_all(text=re.compile('Semiannual Reports'))
    ul = header[0].parent.find_next_sibling()
    for li in ul.find_all('li'):
      link = li.find('a')
      report_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
      report_id = report_url.split('/')[-1].split('.')[0]
      title = 'Semiannual Report to the Congress: %s' % link.text.strip()
      published_on = datetime.datetime.strptime(link.text.strip(), '%B %Y')
      report = self.report_from(report_id, report_url, title, published_on)
      if report:
        inspector.save_report(report)

  def get_bare_reports(self, bare_report_ul):
    for li in bare_report_ul.find_all('li'):
      link = li.find('a')
      report_url = None
      unreleased = False
      if link:
        report_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
        text = link.text
      elif li.find_all(text=RE_PUBLIC_RELEASE):
        unreleased = True
        text = li.get_text()
      else:
        unreleased = True
        text = li.get_text()
        # Annoying reports that aren't public (apparently) but not marked such.
        logging.warning('Found report with no PDF link but not marked "Not for '
                        'Public Release": %s', text)

      title = text.replace(' (Not for Public Release)', '')
      parts = [txt.strip() for txt in text.rsplit(',', 1)]
      if len(parts) == 1:
        published_on = self.date_from_string_or_die(text, report_url)
      else:
        _, published_text = parts
        if unreleased:
          published_text = published_text.replace(
            ' (Not for Public Release)', '')
        published_on = self.date_from_string_or_die(
          published_text, report_url, fallback=text)

      report_id = generate_id(title, published_on)
      report = self.report_from(report_id, report_url, title, published_on)
      if report:
        if unreleased:
          report['unreleased'] = True
          report['landing_url'] = REPORTS_BY_YEAR_URL
        inspector.save_report(report)

  def get_uls_past_audit_header(self, doc):
    # Yes there is mysteriously an extra space in the document before the word
    # 'Reports'.
    header = doc.find_all(text=re.compile('Audits and Other +Reports'))
    return header[0].parent.find_next_siblings('ul')

  def date_from_string_or_die(self, title, report_url, fallback=None):
    md = RE_DATE.search(title)
    if md:
      return datetime.datetime.strptime(md.group(0), '%B %Y')

    md = RE_SHORT_DATE.search(title)
    if md:
      # Seriously IG? Sept?
      return datetime.datetime.strptime(md.group(0).replace('Sept', 'Sep'),
                                        '%b %Y')

    md = RE_FISCAL_YEAR.search(title)
    if md:
      year = int(md.group(1))
      return datetime.datetime(year+1, 11, 1)

    md = RE_LONG_DATE.search(title)
    if md:
      if ',' in md.group(0):
        return datetime.datetime.strptime(md.group(0), '%B %d, %Y')
      else:
        return datetime.datetime.strptime(md.group(0), '%B %d %Y')

    # Ridiculous recursion to avoid exception handling in the top level.
    if fallback:
      try:
        return self.date_from_string_or_die(fallback, report_url)
      except ValueError:
        pass

    # If the string is not a title but instead a published_on_text, this will
    # definitely fail, but it's no big deal to try.
    temp_id = generate_id(title)
    print(temp_id)
    if temp_id in REPORT_ID_TO_EST_DATE:
     return REPORT_ID_TO_EST_DATE[temp_id]

    raise ValueError('Found report with no listed and no hardcoded date:'
                     '%s <%s>' % (title, report_url))

  def report_from(self, report_id, report_url, title, published_on):
    if published_on.year not in self.year_range:
      logging.debug("[%s] Skipping, not in requested range." % report_url)
      return

    agency = 'loc'
    agency_name = 'Library of Congress'
    if RE_OWLC.search(title):
      agency = 'owlc'
      agency_name = 'Open World Leadership Center'
    elif RE_JMCF.search(title):
      agency = 'jmcf'
      agency_name = 'James Madison Council Fund'

    report = {
      'inspector': 'loc',
      'inspector_url': 'http://www.loc.gov/about/office-of-the-inspector-general/',
      'agency': agency,
      'agency_name': agency_name,
      'report_id': report_id,
      'url': report_url,
      'title': title,
      'published_on': datetime.datetime.strftime(published_on, "%Y-%m-%d"),
      'estimated_date': True  # Unfortunately, all of these are estimates.
    }
    return report

def run(options):
  return LibraryOfCongressScraper(options).run()

utils.run(run) if (__name__ == "__main__") else None
<commit_msg>[loc] executable flag<commit_after>#!/usr/bin/env python

import datetime
import logging
import os
from urllib.parse import urljoin
import re

from bs4 import BeautifulSoup
from utils import utils, inspector

# http://www.loc.gov/about/office-of-the-inspector-general/
# Oldest report: 2001

# options:
#   standard since/year options for a year range to fetch from.
#
# Notes for IG's web team:
#

REPORTS_BY_YEAR_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/annual-reports/'
REPORTS_TESTIMONY_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/congressional-testimony/'
REPORTS_PEER_REVIEWS_URL = 'http://www.loc.gov/about/office-of-the-inspector-general/peer-reviews/'

OTHER_REPORTS = {
  'http://lcweb2.loc.gov/master/libn/about/office-of-the-inspector-general/documents/LCR211-6-oct2010.pdf': '2010-10-01'
}

REPORT_ID_TO_EST_DATE = {
  'TLHAOTBMPIAMECT': datetime.datetime(2013, 3, 1),
  'WTTSOTGPARMA': datetime.datetime(2013, 3, 1),
  'MDCITSOICBISCFPR': datetime.datetime(2010, 10, 1),
  'RRSASBRJ2T2A': datetime.datetime(2007, 7, 1),
}

PARTS_MONTH = ('(January|February|March|April|May|June|July|August|September|'
               'October|November|December)')
RE_PUBLIC_RELEASE = re.compile('Not for Public Release', re.I)
RE_DATE = re.compile(r'%s +\d\d\d\d' % PARTS_MONTH)
RE_SHORT_DATE = re.compile(r'(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sept|Oct|Nov|Dec)'
                           r' +\d\d\d\d')
RE_LONG_DATE = re.compile(r'%s \d\d,? +\d\d\d\d' % PARTS_MONTH)
RE_FISCAL_YEAR = re.compile('Fiscal Year (\d\d\d\d)')
RE_OWLC = re.compile('Open +World +Leadership +Center')
RE_JMCF = re.compile('James +Madison +Council +Fund')

def generate_id(title, published_on=None):
  """Generates a report id given the report title and published date.

  This generally tries to be clever and uses the first letter of each name of
  the title, then the month date and year of the report as the ID.
  """
  words = title.split()
  letters = [word[0] for word in words if re.match(r'\w', word[0])]
  slug = ''.join(letters).upper()
  if published_on:
    return '%s-%s' % (slug[:16], published_on.strftime('%Y-%m-%d'))
  else:
    return slug

class LibraryOfCongressScraper(object):
  def __init__(self, options):
    self.options = options
    self.year_range = inspector.year_range(options)

  def run(self):
    self.get_reports_by_year()

    # Get congressional testimony.
    self.get_listed_reports(REPORTS_TESTIMONY_URL)

    # Get peer reviews.
    self.get_listed_reports(REPORTS_PEER_REVIEWS_URL)

  def get_reports_by_year(self):
    # This page contains semianual reports as well as a few Audit reports for
    # Fiscal Year 2014 and links to sub-pages that contain links for other
    # fiscal years.
    doc = BeautifulSoup(utils.download(REPORTS_BY_YEAR_URL))

    # Get the semiannual reports to Congress.
    self.get_semiannual_reports_to_congress(doc)

    # Reports that are 'bare' on the page, listed explicitly
    bare_report_ul, link_subpage_ul = self.get_uls_past_audit_header(doc)
    self.get_bare_reports(bare_report_ul)

    # Links on the page to audit reports from past fiscal years
    for li in link_subpage_ul.find_all('li'):
      link = li.find('a')
      if link:
        next_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
        doc = BeautifulSoup(utils.download(next_url))
        uls = self.get_uls_past_audit_header(doc)
        assert len(uls) == 1, ('Mysterious additional ul data on page: %s' %
                               next_url)
        self.get_bare_reports(uls[0])

  def get_listed_reports(self, url):
    doc = BeautifulSoup(utils.download(url))
    article = doc.select('.article')[0]
    for ul in article.find_all('ul'):
      self.get_bare_reports(ul)

  def get_semiannual_reports_to_congress(self, doc):
    header = doc.find_all(text=re.compile('Semiannual Reports'))
    ul = header[0].parent.find_next_sibling()
    for li in ul.find_all('li'):
      link = li.find('a')
      report_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
      report_id = report_url.split('/')[-1].split('.')[0]
      title = 'Semiannual Report to the Congress: %s' % link.text.strip()
      published_on = datetime.datetime.strptime(link.text.strip(), '%B %Y')
      report = self.report_from(report_id, report_url, title, published_on)
      if report:
        inspector.save_report(report)

  def get_bare_reports(self, bare_report_ul):
    for li in bare_report_ul.find_all('li'):
      link = li.find('a')
      report_url = None
      unreleased = False
      if link:
        report_url = urljoin(REPORTS_BY_YEAR_URL, link['href'])
        text = link.text
      elif li.find_all(text=RE_PUBLIC_RELEASE):
        unreleased = True
        text = li.get_text()
      else:
        unreleased = True
        text = li.get_text()
        # Annoying reports that aren't public (apparently) but not marked such.
        logging.warning('Found report with no PDF link but not marked "Not for '
                        'Public Release": %s', text)

      title = text.replace(' (Not for Public Release)', '')
      parts = [txt.strip() for txt in text.rsplit(',', 1)]
      if len(parts) == 1:
        published_on = self.date_from_string_or_die(text, report_url)
      else:
        _, published_text = parts
        if unreleased:
          published_text = published_text.replace(
            ' (Not for Public Release)', '')
        published_on = self.date_from_string_or_die(
          published_text, report_url, fallback=text)

      report_id = generate_id(title, published_on)
      report = self.report_from(report_id, report_url, title, published_on)
      if report:
        if unreleased:
          report['unreleased'] = True
          report['landing_url'] = REPORTS_BY_YEAR_URL
        inspector.save_report(report)

  def get_uls_past_audit_header(self, doc):
    # Yes there is mysteriously an extra space in the document before the word
    # 'Reports'.
    header = doc.find_all(text=re.compile('Audits and Other +Reports'))
    return header[0].parent.find_next_siblings('ul')

  def date_from_string_or_die(self, title, report_url, fallback=None):
    md = RE_DATE.search(title)
    if md:
      return datetime.datetime.strptime(md.group(0), '%B %Y')

    md = RE_SHORT_DATE.search(title)
    if md:
      # Seriously IG? Sept?
      return datetime.datetime.strptime(md.group(0).replace('Sept', 'Sep'),
                                        '%b %Y')

    md = RE_FISCAL_YEAR.search(title)
    if md:
      year = int(md.group(1))
      return datetime.datetime(year+1, 11, 1)

    md = RE_LONG_DATE.search(title)
    if md:
      if ',' in md.group(0):
        return datetime.datetime.strptime(md.group(0), '%B %d, %Y')
      else:
        return datetime.datetime.strptime(md.group(0), '%B %d %Y')

    # Ridiculous recursion to avoid exception handling in the top level.
    if fallback:
      try:
        return self.date_from_string_or_die(fallback, report_url)
      except ValueError:
        pass

    # If the string is not a title but instead a published_on_text, this will
    # definitely fail, but it's no big deal to try.
    temp_id = generate_id(title)
    print(temp_id)
    if temp_id in REPORT_ID_TO_EST_DATE:
     return REPORT_ID_TO_EST_DATE[temp_id]

    raise ValueError('Found report with no listed and no hardcoded date:'
                     '%s <%s>' % (title, report_url))

  def report_from(self, report_id, report_url, title, published_on):
    if published_on.year not in self.year_range:
      logging.debug("[%s] Skipping, not in requested range." % report_url)
      return

    agency = 'loc'
    agency_name = 'Library of Congress'
    if RE_OWLC.search(title):
      agency = 'owlc'
      agency_name = 'Open World Leadership Center'
    elif RE_JMCF.search(title):
      agency = 'jmcf'
      agency_name = 'James Madison Council Fund'

    report = {
      'inspector': 'loc',
      'inspector_url': 'http://www.loc.gov/about/office-of-the-inspector-general/',
      'agency': agency,
      'agency_name': agency_name,
      'report_id': report_id,
      'url': report_url,
      'title': title,
      'published_on': datetime.datetime.strftime(published_on, "%Y-%m-%d"),
      'estimated_date': True  # Unfortunately, all of these are estimates.
    }
    return report

def run(options):
  return LibraryOfCongressScraper(options).run()

utils.run(run) if (__name__ == "__main__") else None
<|endoftext|>


<commit_before># Intel License for KCM (version January 2017)
#
# Copyright (c) 2017 Intel Corporation.
#
# Use.  You may use the software (the "Software"), without modification,
# provided the following conditions are met:
#
# * Neither the name of Intel nor the names of its suppliers may be used to
#   endorse or promote products derived from this Software without specific
#   prior written permission.
# * No reverse engineering, decompilation, or disassembly of this Software
#   is permitted.
#
# Limited patent license.  Intel grants you a world-wide, royalty-free,
# non-exclusive license under patents it now or hereafter owns or controls to
# make, have made, use, import, offer to sell and sell ("Utilize") this
# Software, but solely to the extent that any such patent is necessary to
# Utilize the Software alone. The patent license shall not apply to any
# combinations which include this software.  No hardware per se is licensed
# hereunder.
#
# Third party and other Intel programs.  "Third Party Programs" are the files
# listed in the "third-party-programs.txt" text file that is included with the
# Software and may include Intel programs under separate license terms. Third
# Party Programs, even if included with the distribution of the Materials, are
# governed by separate license terms and those license terms solely govern your
# use of those programs.
#
# DISCLAIMER.  THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT ARE
# DISCLAIMED. THIS SOFTWARE IS NOT INTENDED NOR AUTHORIZED FOR USE IN SYSTEMS
# OR APPLICATIONS WHERE FAILURE OF THE SOFTWARE MAY CAUSE PERSONAL INJURY OR
# DEATH.
#
# LIMITATION OF LIABILITY. IN NO EVENT WILL INTEL BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. YOU AGREE TO
# INDEMNIFIY AND HOLD INTEL HARMLESS AGAINST ANY CLAIMS AND EXPENSES RESULTING
# FROM YOUR USE OR UNAUTHORIZED USE OF THE SOFTWARE.
#
# No support.  Intel may make changes to the Software, at any time without
# notice, and is not obligated to support, update or provide training for the
# Software.
#
# Termination. Intel may terminate your right to use the Software in the event
# of your breach of this Agreement and you fail to cure the breach within a
# reasonable period of time.
#
# Feedback.  Should you provide Intel with comments, modifications,
# corrections, enhancements or other input ("Feedback") related to the Software
# Intel will be free to use, disclose, reproduce, license or otherwise
# distribute or exploit the Feedback in its sole discretion without any
# obligations or restrictions of any kind, including without limitation,
# intellectual property rights or licensing obligations.
#
# Compliance with laws.  You agree to comply with all relevant laws and
# regulations governing your use, transfer, import or export (or prohibition
# thereof) of the Software.
#
# Governing law.  All disputes will be governed by the laws of the United
# States of America and the State of Delaware without reference to conflict of
# law principles and subject to the exclusive jurisdiction of the state or
# federal courts sitting in the State of Delaware, and each party agrees that
# it submits to the personal jurisdiction and venue of those courts and waives
# any objections. The United Nations Convention on Contracts for the
# International Sale of Goods (1980) is specifically excluded and will not
# apply to the Software.

import ast
from intel import config
import json
from kubernetes import config as k8sconfig, client as k8sclient
from kubernetes.client.rest import ApiException as K8sApiException
import logging
import os
import sys

# discover reads the KCM configuration file, patches kubernetes nodes with
# appropriate number of KCM Opaque Integer Resource (OIR) slots and applies
# the appropriate KCM node labels and taints.
def discover(conf_dir):
    # Patch the node with the appropriate KCM OIR.
    logging.debug("Patching the node with the appropriate KCM OIR.")
    add_node_oir(conf_dir)
    # Add appropriate KCM label to the node.
    logging.debug("Adding appropriate KCM label to the node.")
    add_node_label()
    # Add appropriate KCM taint to the node.
    logging.debug("Adding appropriate KCM taint to the node.")
    add_node_taint()

# add_node_oir patches the node with the appropriate KCM OIR.
def add_node_oir(conf_dir):
    c = config.Config(conf_dir)
    with c.lock():
        if "dataplane" not in c.pools():
            raise KeyError("Dataplane pool does not exist")
        if len(c.pool("dataplane").cpu_lists()) == 0:
            raise KeyError("No CPU list in dataplane pool")
        num_slots = len(c.pool("dataplane").cpu_lists())

    patch_path = ('/status/capacity/pod.alpha.kubernetes.io~1opaque-int-'
                  'resource-kcm')
    patch_body = [{
        "op": "add",
        "path": patch_path,
        "value": num_slots
    }]

    try:
        patch_k8s_node_status(patch_body)
    except K8sApiException as err:
        logging.error("Exception when patching node with OIR: {}"
                      .format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

def add_node_label():
    patch_path = '/metadata/labels/kcm.intel.com~1kcm-node'
    patch_body = [{
        "op": "add",
        "path": patch_path,
        "value": "true"
    }]

    try:
        patch_k8s_node(patch_body)
    except K8sApiException as err:
        logging.error("Exception when labeling the node: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

def add_node_taint():
    node_name = os.getenv("NODE_NAME")
    try:
        node_resp = get_k8s_node(node_name)
    except K8sApiException as err:
        logging.error("Exception when getting the node obj: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

    node_taints_list = []
    node_taint_key = "scheduler.alpha.kubernetes.io/taints"
    if node_taint_key in node_resp["metadata"]["annotations"]:
        node_taints = node_resp["metadata"]["annotations"][node_taint_key]
        node_taints_list = ast.literal_eval(node_taints)

    kcm_taint = {
            "key": "kcm",
            "value": "true",
            "effect": "NoSchedule"
    }
    node_taints_list.append(kcm_taint)

    patch_path = '/metadata/annotations/scheduler.alpha.kubernetes.io~1taints'
    patch_body = [{
        "op": "replace",
        "path": patch_path,
        "value": json.dumps(node_taints_list)
    }]

    try:
        patch_k8s_node(patch_body)
    except K8sApiException as err:
        logging.error("Exception when tainting the node: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

# patch_k8s_node_status patches the kubernetes node with KCM OIR with
# value num_slots.
def patch_k8s_node_status(patch_body):
    k8sconfig.load_incluster_config()
    k8sapi = k8sclient.CoreV1Api()
    node_name = os.getenv("NODE_NAME")

    # Patch the node with the specified number of opaque integer resources.
    k8sapi.patch_node_status(node_name, patch_body)

def patch_k8s_node(patch_body):
    k8sconfig.load_incluster_config()
    k8sapi = k8sclient.CoreV1Api()
    node_name = os.getenv("NODE_NAME")

    # Patch the node with the specified number of opaque integer resources.
    k8sapi.patch_node(node_name, patch_body)

# get_k8s_node() retuns the node spec associated with node_name.
def get_k8s_node(node_name):
    node_list_resp = get_k8s_node_list()
    for node in node_list_resp["items"]:
        if node["metadata"]["name"] == node_name:
            return node

# get_k8s_node_list() returns the node list in the current Kubernetes cluster.
def get_k8s_node_list():
    k8sconfig.load_incluster_config()
    k8s_api = k8sclient.CoreV1Api()
    return k8s_api.list_node().to_dict()
<commit_msg>`kcm discover` does not republish existing taint.<commit_after># Intel License for KCM (version January 2017)
#
# Copyright (c) 2017 Intel Corporation.
#
# Use.  You may use the software (the "Software"), without modification,
# provided the following conditions are met:
#
# * Neither the name of Intel nor the names of its suppliers may be used to
#   endorse or promote products derived from this Software without specific
#   prior written permission.
# * No reverse engineering, decompilation, or disassembly of this Software
#   is permitted.
#
# Limited patent license.  Intel grants you a world-wide, royalty-free,
# non-exclusive license under patents it now or hereafter owns or controls to
# make, have made, use, import, offer to sell and sell ("Utilize") this
# Software, but solely to the extent that any such patent is necessary to
# Utilize the Software alone. The patent license shall not apply to any
# combinations which include this software.  No hardware per se is licensed
# hereunder.
#
# Third party and other Intel programs.  "Third Party Programs" are the files
# listed in the "third-party-programs.txt" text file that is included with the
# Software and may include Intel programs under separate license terms. Third
# Party Programs, even if included with the distribution of the Materials, are
# governed by separate license terms and those license terms solely govern your
# use of those programs.
#
# DISCLAIMER.  THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT ARE
# DISCLAIMED. THIS SOFTWARE IS NOT INTENDED NOR AUTHORIZED FOR USE IN SYSTEMS
# OR APPLICATIONS WHERE FAILURE OF THE SOFTWARE MAY CAUSE PERSONAL INJURY OR
# DEATH.
#
# LIMITATION OF LIABILITY. IN NO EVENT WILL INTEL BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. YOU AGREE TO
# INDEMNIFIY AND HOLD INTEL HARMLESS AGAINST ANY CLAIMS AND EXPENSES RESULTING
# FROM YOUR USE OR UNAUTHORIZED USE OF THE SOFTWARE.
#
# No support.  Intel may make changes to the Software, at any time without
# notice, and is not obligated to support, update or provide training for the
# Software.
#
# Termination. Intel may terminate your right to use the Software in the event
# of your breach of this Agreement and you fail to cure the breach within a
# reasonable period of time.
#
# Feedback.  Should you provide Intel with comments, modifications,
# corrections, enhancements or other input ("Feedback") related to the Software
# Intel will be free to use, disclose, reproduce, license or otherwise
# distribute or exploit the Feedback in its sole discretion without any
# obligations or restrictions of any kind, including without limitation,
# intellectual property rights or licensing obligations.
#
# Compliance with laws.  You agree to comply with all relevant laws and
# regulations governing your use, transfer, import or export (or prohibition
# thereof) of the Software.
#
# Governing law.  All disputes will be governed by the laws of the United
# States of America and the State of Delaware without reference to conflict of
# law principles and subject to the exclusive jurisdiction of the state or
# federal courts sitting in the State of Delaware, and each party agrees that
# it submits to the personal jurisdiction and venue of those courts and waives
# any objections. The United Nations Convention on Contracts for the
# International Sale of Goods (1980) is specifically excluded and will not
# apply to the Software.

from intel import config
import json
from kubernetes import config as k8sconfig, client as k8sclient
from kubernetes.client.rest import ApiException as K8sApiException
import logging
import os
import sys

# discover reads the KCM configuration file, patches kubernetes nodes with
# appropriate number of KCM Opaque Integer Resource (OIR) slots and applies
# the appropriate KCM node labels and taints.
def discover(conf_dir):
    # Patch the node with the appropriate KCM OIR.
    logging.debug("Patching the node with the appropriate KCM OIR.")
    add_node_oir(conf_dir)
    # Add appropriate KCM label to the node.
    logging.debug("Adding appropriate KCM label to the node.")
    add_node_label()
    # Add appropriate KCM taint to the node.
    logging.debug("Adding appropriate KCM taint to the node.")
    add_node_taint()

# add_node_oir patches the node with the appropriate KCM OIR.
def add_node_oir(conf_dir):
    c = config.Config(conf_dir)
    with c.lock():
        if "dataplane" not in c.pools():
            raise KeyError("Dataplane pool does not exist")
        if len(c.pool("dataplane").cpu_lists()) == 0:
            raise KeyError("No CPU list in dataplane pool")
        num_slots = len(c.pool("dataplane").cpu_lists())

    patch_path = ("/status/capacity/pod.alpha.kubernetes.io~1opaque-int-"
                  "resource-kcm")
    patch_body = [{
        "op": "add",
        "path": patch_path,
        "value": num_slots
    }]

    try:
        patch_k8s_node_status(patch_body)
    except K8sApiException as err:
        logging.error("Exception when patching node with OIR: {}"
                      .format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

def add_node_label():
    patch_path = "/metadata/labels/kcm.intel.com~1kcm-node"
    patch_body = [{
        "op": "add",
        "path": patch_path,
        "value": "true"
    }]

    try:
        patch_k8s_node(patch_body)
    except K8sApiException as err:
        logging.error("Exception when labeling the node: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

def add_node_taint():
    node_name = os.getenv("NODE_NAME")
    try:
        node_resp = get_k8s_node(node_name)
    except K8sApiException as err:
        logging.error("Exception when getting the node obj: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

    node_taints_list = []
    node_taint_key = "scheduler.alpha.kubernetes.io/taints"
    if node_taint_key in node_resp["metadata"]["annotations"]:
        node_taints = node_resp["metadata"]["annotations"][node_taint_key]
        # Do not try to parse the empty string as JSON.
        if node_taints:
            node_taints_list = json.loads(node_taints)

    # Filter existing "kcm" taint, if it exists.
    node_taints_list = [t for t in node_taints_list if t["key"] != "kcm"]

    node_taints_list.append({
        "key": "kcm",
        "value": "true",
        "effect": "NoSchedule"
    })

    patch_path = "/metadata/annotations/scheduler.alpha.kubernetes.io~1taints"

    # See: https://tools.ietf.org/html/rfc6902#section-4.1
    patch_body = [
        {
            "op": "add",
            "path": patch_path,
            "value": json.dumps(node_taints_list)
        }
    ]

    try:
        patch_k8s_node(patch_body)
    except K8sApiException as err:
        logging.error("Exception when tainting the node: {}".format(err))
        logging.error("Aborting discover ...")
        sys.exit(1)

# patch_k8s_node_status patches the kubernetes node with KCM OIR with
# value num_slots.
def patch_k8s_node_status(patch_body):
    k8sconfig.load_incluster_config()
    k8sapi = k8sclient.CoreV1Api()
    node_name = os.getenv("NODE_NAME")

    logging.info("Patching node status {}:\n{}".format(
        node_name,
        json.dumps(patch_body, indent=2, sort_keys=True)))

    # Patch the node with the specified number of opaque integer resources.
    k8sapi.patch_node_status(node_name, patch_body)

def patch_k8s_node(patch_body):
    k8sconfig.load_incluster_config()
    k8sapi = k8sclient.CoreV1Api()
    node_name = os.getenv("NODE_NAME")

    logging.info("Patching node {}:\n{}".format(
        node_name,
        json.dumps(patch_body, indent=2, sort_keys=True)))

    # Patch the node with the specified number of opaque integer resources.
    k8sapi.patch_node(node_name, patch_body)

# get_k8s_node() retuns the node spec associated with node_name.
def get_k8s_node(node_name):
    node_list_resp = get_k8s_node_list()
    for node in node_list_resp["items"]:
        if node["metadata"]["name"] == node_name:
            return node

# get_k8s_node_list() returns the node list in the current Kubernetes cluster.
def get_k8s_node_list():
    k8sconfig.load_incluster_config()
    k8s_api = k8sclient.CoreV1Api()
    return k8s_api.list_node().to_dict()
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function
from multipledispatch import dispatch

from blaze.expr.table import *

__all__ = ['compute']

base = (int, float, str, bool)

@dispatch(base, object)
def compute(a, b):
    return a

@dispatch(TableExpr, dict)
def compute(t, d):
    if t in d:
        return d[t]

    # Pare down d to only nodes in table
    nodes = set(t.traverse())
    d = dict((k, v) for k, v in d.items()
                    if k in nodes)

    # Common case: One relevant value in dict
    #              Switch to standard dispatching scheme
    if len(d) == 1:
        return compute(t, list(d.values())[0])

    if hasattr(t, 'parent'):
        parent = compute(t.parent, d)
        t2 = t.subs({t.parent: TableSymbol('', t.parent.schema)})
        return compute(t2, parent)

    raise NotImplementedError("No method found to compute on multiple Tables")

@dispatch(Join, dict)
def compute(t, d):
    lhs = compute(t.lhs, d)
    rhs = compute(t.rhs, d)

    t2 = t.subs({t.lhs: TableSymbol('_lhs', t.lhs.schema),
                 t.rhs: TableSymbol('_rhs', t.rhs.schema)})
    return compute(t2, lhs, rhs)

@dispatch(Join, object)
def compute(t, o):
    return compute(t, o, o)

def columnwise_funcstr(t):
    """
    >>> t = TableSymbol('t', '{x: real, y: real}')
    >>> cw = t['x'] + t['y']
    >>> columnwise_funcstr(cw)
    'lambda a, b: a + b'
    """
    return 'lambda %s: %s' % (', '.join(map(eval_str, t.argsymbols)),
                              eval_str(t.expr))
<commit_msg>columnwise_funcstr provides full and variadic versions<commit_after>from __future__ import absolute_import, division, print_function
from multipledispatch import dispatch

from blaze.expr.table import *

__all__ = ['compute']

base = (int, float, str, bool)

@dispatch(base, object)
def compute(a, b):
    return a

@dispatch(TableExpr, dict)
def compute(t, d):
    if t in d:
        return d[t]

    # Pare down d to only nodes in table
    nodes = set(t.traverse())
    d = dict((k, v) for k, v in d.items()
                    if k in nodes)

    # Common case: One relevant value in dict
    #              Switch to standard dispatching scheme
    if len(d) == 1:
        return compute(t, list(d.values())[0])

    if hasattr(t, 'parent'):
        parent = compute(t.parent, d)
        t2 = t.subs({t.parent: TableSymbol('', t.parent.schema)})
        return compute(t2, parent)

    raise NotImplementedError("No method found to compute on multiple Tables")

@dispatch(Join, dict)
def compute(t, d):
    lhs = compute(t.lhs, d)
    rhs = compute(t.rhs, d)

    t2 = t.subs({t.lhs: TableSymbol('_lhs', t.lhs.schema),
                 t.rhs: TableSymbol('_rhs', t.rhs.schema)})
    return compute(t2, lhs, rhs)

@dispatch(Join, object)
def compute(t, o):
    return compute(t, o, o)

def columnwise_funcstr(t, variadic=True, full=False):
    """
    >>> t = TableSymbol('t', '{x: real, y: real, z: real}')
    >>> cw = t['x'] + t['z']
    >>> columnwise_funcstr(cw)
    'lambda x, z: x + z'

    >>> columnwise_funcstr(cw, variadic=False)
    'lambda (x, z): x + z'

    >>> columnwise_funcstr(cw, variadic=False, full=True)
    'lambda (x, y, z): x + z'
    """
    if full:
        columns = t.parent.columns
    else:
        columns = t.active_columns()
    if variadic:
        prefix = 'lambda %s: '
    else:
        prefix = 'lambda (%s): '

    return prefix % ', '.join(map(str, columns)) + eval_str(t.expr)
<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from __future__ import absolute_import, division, print_function
from future.utils.six import StringIO

import unittest

from skbio import TreeNode
from skbio.io import NewickFormatError
from skbio.io.newick import (_newick_to_tree_node, _tree_node_to_newick,
                             _newick_sniffer)

class TestNewick(unittest.TestCase):
    def _invert_tree(self, tree):
        tree.children = [self._invert_tree(child) for child
                         in reversed(tree.children)]
        return tree

    def _is_node_equal(self, n1, n2):
        self.assertEqual(n1.name, n2.name)
        self.assertEqual(n1.length, n2.length)
        self.assertEqual(len(n1.children), len(n2.children))
        return True

    def _is_equal(self, n1, n2):
        name = lambda x: (x.name, x.length, len(x.children))
        if self._is_node_equal(n1, n2):
            for c1, c2 in zip(sorted(n1.children if n1.children else [],
                                     key=name),
                              sorted(n2.children if n1.children else [],
                                     key=name)):
                self.assertTrue(c1.parent is n1)
                self.assertTrue(c2.parent is n2)
                if not self._is_equal(c1, c2):
                    return False
            return True
        return False

    def _setup_tree(self, kwargs_list):
        trees = []
        for kwargs in kwargs_list:
            trees.append(TreeNode(**kwargs))

        trees[4].extend([trees[2], trees[3]])
        trees[5].extend([trees[0], trees[1], trees[4]])

        return trees[5]

    def _setup_linked_list(self, kwargs_list):
        last_node = None
        for idx, kwargs in enumerate(kwargs_list):
            new_node = TreeNode(**kwargs)

            if last_node is not None:
                new_node.append(last_node)
            last_node = new_node
        return last_node

    def _setup_balanced_binary(self, kwargs_list):
        trees = []
        for kwargs in kwargs_list:
            trees.append(TreeNode(**kwargs))

        trees[0].extend([trees[2], trees[3]])
        trees[1].extend([trees[4], trees[5]])
        trees[6].extend([trees[0], trees[1]])
        return trees[6]

    def setUp(self):
        tree_blank = (self._setup_tree([
            {}, {}, {}, {}, {}, {}
        ]), [
            "(,,(,));",
            "(,(,),);",
            "((,),,);",
            "   ((,[ this is a comment ])      ,    ,   )    ;  ",
            "((,[ i_can_do_this[0] or escape unmatched '[ ]),[more words],);",
        ])

        tree_leaves_named = (self._setup_tree([
            {'name': 'a_'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': 'd'},
            {},
            {}
        ]), [
            "('a_',b,(c,d));",
            "(b,(c,d),'a_');",
            "(b\n,'a_'\n  ,(d \t,c) )  ;",
        ])

        tree_all_named = (self._setup_tree([
            {'name': 'a'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': '[whaaat!\']'},
            {'name': 'e'},
            {'name': 'f'}
        ]), [
            "(a,b,(c,'[whaaat!'']')e)f;",
            "(b,(c,'[whaaat!'']')e,a)f;",
            "(b,[comment] \na,('[whaaat!'']',c)e)f;",
        ])

        tree_all_but_root_distances = (self._setup_tree([
            {'length': 0.1},
            {'length': 0.2},
            {'length': 0.3},
            {'length': 0.4},
            {'length': 0.5},
            {}
        ]), [
            "(:0.1,:0.2,(:0.3,:0.4):0.5);",
            "(:0.2,(:0.3,:0.4):0.5,:0.1);",
            "(:0.2,:0.1,(:0.4,:0.3):0.5);",
        ])

        tree_all_distances = (self._setup_tree([
            {'length': 0.1},
            {'length': 0.2},
            {'length': 0.3},
            {'length': 0.4},
            {'length': 0.5},
            {'length': 0.0}
        ]), [
            "(:0.1,:0.2,(:0.3,:0.4):0.5):0.0;",
            "(:0.2,(:0.3,:0.4):0.5,:0.1):0.0;",
            "(:0.2,\n:0.1,(:0.4,\n:0.3):0.5)\n:0.0;",
        ])

        tree_all_leaves_named_with_distances = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de d', 'length': 0.4},
            {'length': 0.5},
            {'length': 0.0}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de_d:0.4):0.5):0.0;",
            "('b_a''':0.2,(c:0.3,'de d':0.4):0.5,a:0.1):0.0;",
            "('b_a''':0.2,a:0.1,('de d'[why not]:0.4,c:0.3):0.5):0.0;",
        ])

        tree_all_leaves_named_with_distances_no_root = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de  d', 'length': 0.4},
            {'length': 0.5},
            {}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de__d:0.4):0.5);",
            "('b_a''':0.2\n[comment ahoy]\n,(c:0.3,'de  d':0.4):0.5,a:0.1);",
            "('b_a''':0.2,a:0.1,(de__d:0.4,c:0.3):0.5);"
        ])

        tree_all = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de\' d', 'length': 0.4},
            {'name': 'e', 'length': 0.5},
            {'name': 'f', 'length': 0.0}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de''_d:0.4)e:0.5)f:0.0;",
            "('b_a''':0.2,(c:0.3,de''_d:0.4)e:0.5,a:0.1)f:0.0;",
            "((de''_d:0.4, c:0.3)e:0.5, 'b_a''':0.2, a:0.1)f:0.0;"
        ])

        balanced_blank = (self._setup_balanced_binary([
            {}, {}, {}, {}, {}, {}, {}
        ]), [
            "((,),(,));"
        ])

        balanced_named = (self._setup_balanced_binary([
            {'name': 'a'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': 'd'},
            {'name': 'e'},
            {'name': 'f'},
            {'name': 'g'}
        ]), [
            "((c,d)a,(e,f)b)g;"
        ])

        balanced_distances = (self._setup_balanced_binary([
            {'length': 1.0},
            {'length': 2.0},
            {'length': 3.0},
            {'length': 4.0},
            {'length': 5.0},
            {'length': 6.0},
            {'length': 0.0}
        ]), [
            "((:3.0,:4.0):1.0,(:5.0,:6.0):2.0):0.0;"
        ])

        blanaced_all = (self._setup_balanced_binary([
            {'name': 'a', 'length': 1.0},
            {'name': 'b', 'length': 2.0},
            {'name': 'c', 'length': 3.0},
            {'name': 'd', 'length': 4.0},
            {'name': 'e', 'length': 5.0},
            {'name': 'f:f\'f', 'length': 6.0},
            {'name': 'g', 'length': 0.0}
        ]), [
            "((c:3.0,d:4.0)a:1.0,(e:5.0,'f:f''f':6.0)b:2.0)g:0.0;"
        ])

        linked_list_blank = (self._setup_linked_list([
            {}, {}, {}, {}, {}
        ]), [
            "(((())));",
            "[(((())));](((())));",
            "[[(((())));](((())));](((())));\t\t\n"
        ])

        linked_list_named = (self._setup_linked_list([
            {'name': 'aaa'},
            {'name': 'b_a\''},
            {'name': 'c'},
            {'name': 'de d'},
            {'name': 'e'},
        ]), [
            "((((aaa)'b_a''')c)de_d)e;"
        ])

        inked_list_distances = (self._setup_linked_list([
            {'length': 0.4},
            {'length': 0.3},
            {'length': 0.2},
            {'length': 0.1},
            {'length': 0.0},
        ]), [
            "((((:0.4):0.3):0.2):0.1):0.0;",
            "((((:0.4)[not a label]:0.3):0.2):0.1):0.0;\t\t\n"
        ])

        linked_list_all = (self._setup_linked_list([
            {'name': 'a', 'length': 0.4},
            {'name': 'b_a\'', 'length': 0.3},
            {'name': 'c', 'length': 0.2},
            {'name': 'de d', 'length': 0.1},
            {'name': 'eee', 'length': 0.0},
        ]), [
            "((((a:0.4)'b_a''':0.3)c:0.2)de_d:0.1)eee:0.0;"
        ])

        single_empty = (TreeNode(), [";", "[comment about the root"
                                     " and it's properties];"])
        single_named = (TreeNode(name='athing'), ["athing;"])
        single_distance = (TreeNode(length=200.0), [":200.0;"])
        single_all = (TreeNode(name='[a]', length=200.0), ["'[a]':200.0;"])

        self.trees_newick_lists = [
            tree_blank,
            tree_leaves_named,
            tree_all_named,
            tree_all_but_root_distances,
            tree_all_distances,
            tree_all_leaves_named_with_distances,
            tree_all_leaves_named_with_distances_no_root,
            tree_all,
            balanced_blank,
            balanced_named,
            balanced_distances,
            blanaced_all,
            linked_list_blank,
            linked_list_named,
            inked_list_distances,
            linked_list_all,
            single_empty,
            single_named,
            single_distance,
            single_all
        ]

        self.invalid_newicks = [
            "",
            "This is not a newick file.",
            "((();",
            "(,,,)(,);",
            "(()());",
            "(():,,)",
            "[][[]('comment is the gotcha':0.2,,);",
            "#SampleID\tHeaderA\tHeaderB\n0\t'yellow'\t0.45;",
            "))();",
            "((,,),((,,));"
        ]

    def test_newick_to_tree_node_valid_files(self):
        for tree, newicks in self.trees_newick_lists:
            for newick in newicks:
                fh = StringIO(newick)
                read_tree = _newick_to_tree_node(fh)

                self.assertTrue(self._is_equal(tree, read_tree))

                fh.close()

    def test_newick_to_tree_node_invalid_files(self):
        for invalid in self.invalid_newicks:
            fh = StringIO(invalid)
            with self.assertRaises(NewickFormatError):
                _newick_to_tree_node(fh)
            fh.close()

    def test_tree_node_to_newick(self):
        for tree, newicks in self.trees_newick_lists:
            newick = newicks[0]
            fh = StringIO()
            _tree_node_to_newick(tree, fh)

            self.assertEqual(newick, fh.getvalue())

            fh.close()

    def test_roundtrip(self):
        for tree, newicks in self.trees_newick_lists:
            newick = newicks[0]
            fh = StringIO(newick)
            tree = _newick_to_tree_node(fh)
            fh2 = StringIO()
            _tree_node_to_newick(tree, fh2)
            fh2.seek(0)
            tree2 = _newick_to_tree_node(fh2)

            self.assertEqual(newick, fh2.getvalue())
            self.assertTrue(self._is_equal(tree, tree2))

            fh.close()
            fh2.close()

    def test_newick_to_tree_node_convert_underscores(self):
        fh = StringIO('(_:0.1, _a, _b)__;')
        tree = _newick_to_tree_node(fh, convert_underscores=False)
        fh2 = StringIO()
        _tree_node_to_newick(tree, fh2)
        self.assertEquals(fh2.getvalue(), "('_':0.1,'_a','_b')'__';")
        fh2.close()
        fh.close()

    def test_newick_sniffer_valid_files(self):
        for _, newicks in self.trees_newick_lists:
            for newick in newicks:
                fh = StringIO(newick)
                self.assertEqual(_newick_sniffer(fh), (True, {}))
                fh.close()

    def test_newick_sniffer_invalid_files(self):
        for invalid in self.invalid_newicks:
            fh = StringIO(invalid)
            self.assertEqual(_newick_sniffer(fh), (False, {}))
            fh.close()

if __name__ == '__main__':
    unittest.main()
<commit_msg>TST: py3 - this better work...<commit_after># ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from __future__ import absolute_import, division, print_function
from future.utils.six import StringIO

import unittest

from skbio import TreeNode
from skbio.io import NewickFormatError
from skbio.io.newick import (_newick_to_tree_node, _tree_node_to_newick,
                             _newick_sniffer)

class TestNewick(unittest.TestCase):
    def _invert_tree(self, tree):
        tree.children = [self._invert_tree(child) for child
                         in reversed(tree.children)]
        return tree

    def _is_node_equal(self, n1, n2):
        self.assertEqual(n1.name, n2.name)
        self.assertEqual(n1.length, n2.length)
        self.assertEqual(len(n1.children), len(n2.children))
        return True

    def _is_equal(self, n1, n2):
        name = lambda x: (str(x.name), float(x.length), len(x.children))
        if self._is_node_equal(n1, n2):
            for c1, c2 in zip(sorted(n1.children, key=name),
                              sorted(n2.children, key=name)):
                self.assertTrue(c1.parent is n1)
                self.assertTrue(c2.parent is n2)
                if not self._is_equal(c1, c2):
                    return False
            return True
        return False

    def _setup_tree(self, kwargs_list):
        trees = []
        for kwargs in kwargs_list:
            trees.append(TreeNode(**kwargs))

        trees[4].extend([trees[2], trees[3]])
        trees[5].extend([trees[0], trees[1], trees[4]])

        return trees[5]

    def _setup_linked_list(self, kwargs_list):
        last_node = None
        for idx, kwargs in enumerate(kwargs_list):
            new_node = TreeNode(**kwargs)

            if last_node is not None:
                new_node.append(last_node)
            last_node = new_node
        return last_node

    def _setup_balanced_binary(self, kwargs_list):
        trees = []
        for kwargs in kwargs_list:
            trees.append(TreeNode(**kwargs))

        trees[0].extend([trees[2], trees[3]])
        trees[1].extend([trees[4], trees[5]])
        trees[6].extend([trees[0], trees[1]])
        return trees[6]

    def setUp(self):
        tree_blank = (self._setup_tree([
            {}, {}, {}, {}, {}, {}
        ]), [
            "(,,(,));",
            "(,(,),);",
            "((,),,);",
            "   ((,[ this is a comment ])      ,    ,   )    ;  ",
            "((,[ i_can_do_this[0] or escape unmatched '[ ]),[more words],);",
        ])

        tree_leaves_named = (self._setup_tree([
            {'name': 'a_'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': 'd'},
            {},
            {}
        ]), [
            "('a_',b,(c,d));",
            "(b,(c,d),'a_');",
            "(b\n,'a_'\n  ,(d \t,c) )  ;",
        ])

        tree_all_named = (self._setup_tree([
            {'name': 'a'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': '[whaaat!\']'},
            {'name': 'e'},
            {'name': 'f'}
        ]), [
            "(a,b,(c,'[whaaat!'']')e)f;",
            "(b,(c,'[whaaat!'']')e,a)f;",
            "(b,[comment] \na,('[whaaat!'']',c)e)f;",
        ])

        tree_all_but_root_distances = (self._setup_tree([
            {'length': 0.1},
            {'length': 0.2},
            {'length': 0.3},
            {'length': 0.4},
            {'length': 0.5},
            {}
        ]), [
            "(:0.1,:0.2,(:0.3,:0.4):0.5);",
            "(:0.2,(:0.3,:0.4):0.5,:0.1);",
            "(:0.2,:0.1,(:0.4,:0.3):0.5);",
        ])

        tree_all_distances = (self._setup_tree([
            {'length': 0.1},
            {'length': 0.2},
            {'length': 0.3},
            {'length': 0.4},
            {'length': 0.5},
            {'length': 0.0}
        ]), [
            "(:0.1,:0.2,(:0.3,:0.4):0.5):0.0;",
            "(:0.2,(:0.3,:0.4):0.5,:0.1):0.0;",
            "(:0.2,\n:0.1,(:0.4,\n:0.3):0.5)\n:0.0;",
        ])

        tree_all_leaves_named_with_distances = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de d', 'length': 0.4},
            {'length': 0.5},
            {'length': 0.0}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de_d:0.4):0.5):0.0;",
            "('b_a''':0.2,(c:0.3,'de d':0.4):0.5,a:0.1):0.0;",
            "('b_a''':0.2,a:0.1,('de d'[why not]:0.4,c:0.3):0.5):0.0;",
        ])

        tree_all_leaves_named_with_distances_no_root = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de  d', 'length': 0.4},
            {'length': 0.5},
            {}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de__d:0.4):0.5);",
            "('b_a''':0.2\n[comment ahoy]\n,(c:0.3,'de  d':0.4):0.5,a:0.1);",
            "('b_a''':0.2,a:0.1,(de__d:0.4,c:0.3):0.5);"
        ])

        tree_all = (self._setup_tree([
            {'name': 'a', 'length': 0.1},
            {'name': 'b_a\'', 'length': 0.2},
            {'name': 'c', 'length': 0.3},
            {'name': 'de\' d', 'length': 0.4},
            {'name': 'e', 'length': 0.5},
            {'name': 'f', 'length': 0.0}
        ]), [
            "(a:0.1,'b_a''':0.2,(c:0.3,de''_d:0.4)e:0.5)f:0.0;",
            "('b_a''':0.2,(c:0.3,de''_d:0.4)e:0.5,a:0.1)f:0.0;",
            "((de''_d:0.4, c:0.3)e:0.5, 'b_a''':0.2, a:0.1)f:0.0;"
        ])

        balanced_blank = (self._setup_balanced_binary([
            {}, {}, {}, {}, {}, {}, {}
        ]), [
            "((,),(,));"
        ])

        balanced_named = (self._setup_balanced_binary([
            {'name': 'a'},
            {'name': 'b'},
            {'name': 'c'},
            {'name': 'd'},
            {'name': 'e'},
            {'name': 'f'},
            {'name': 'g'}
        ]), [
            "((c,d)a,(e,f)b)g;"
        ])

        balanced_distances = (self._setup_balanced_binary([
            {'length': 1.0},
            {'length': 2.0},
            {'length': 3.0},
            {'length': 4.0},
            {'length': 5.0},
            {'length': 6.0},
            {'length': 0.0}
        ]), [
            "((:3.0,:4.0):1.0,(:5.0,:6.0):2.0):0.0;"
        ])

        blanaced_all = (self._setup_balanced_binary([
            {'name': 'a', 'length': 1.0},
            {'name': 'b', 'length': 2.0},
            {'name': 'c', 'length': 3.0},
            {'name': 'd', 'length': 4.0},
            {'name': 'e', 'length': 5.0},
            {'name': 'f:f\'f', 'length': 6.0},
            {'name': 'g', 'length': 0.0}
        ]), [
            "((c:3.0,d:4.0)a:1.0,(e:5.0,'f:f''f':6.0)b:2.0)g:0.0;"
        ])

        linked_list_blank = (self._setup_linked_list([
            {}, {}, {}, {}, {}
        ]), [
            "(((())));",
            "[(((())));](((())));",
            "[[(((())));](((())));](((())));\t\t\n"
        ])

        linked_list_named = (self._setup_linked_list([
            {'name': 'aaa'},
            {'name': 'b_a\''},
            {'name': 'c'},
            {'name': 'de d'},
            {'name': 'e'},
        ]), [
            "((((aaa)'b_a''')c)de_d)e;"
        ])

        inked_list_distances = (self._setup_linked_list([
            {'length': 0.4},
            {'length': 0.3},
            {'length': 0.2},
            {'length': 0.1},
            {'length': 0.0},
        ]), [
            "((((:0.4):0.3):0.2):0.1):0.0;",
            "((((:0.4)[not a label]:0.3):0.2):0.1):0.0;\t\t\n"
        ])

        linked_list_all = (self._setup_linked_list([
            {'name': 'a', 'length': 0.4},
            {'name': 'b_a\'', 'length': 0.3},
            {'name': 'c', 'length': 0.2},
            {'name': 'de d', 'length': 0.1},
            {'name': 'eee', 'length': 0.0},
        ]), [
            "((((a:0.4)'b_a''':0.3)c:0.2)de_d:0.1)eee:0.0;"
        ])

        single_empty = (TreeNode(), [";", "[comment about the root"
                                     " and it's properties];"])
        single_named = (TreeNode(name='athing'), ["athing;"])
        single_distance = (TreeNode(length=200.0), [":200.0;"])
        single_all = (TreeNode(name='[a]', length=200.0), ["'[a]':200.0;"])

        self.trees_newick_lists = [
            tree_blank,
            tree_leaves_named,
            tree_all_named,
            tree_all_but_root_distances,
            tree_all_distances,
            tree_all_leaves_named_with_distances,
            tree_all_leaves_named_with_distances_no_root,
            tree_all,
            balanced_blank,
            balanced_named,
            balanced_distances,
            blanaced_all,
            linked_list_blank,
            linked_list_named,
            inked_list_distances,
            linked_list_all,
            single_empty,
            single_named,
            single_distance,
            single_all
        ]

        self.invalid_newicks = [
            "",
            "This is not a newick file.",
            "((();",
            "(,,,)(,);",
            "(()());",
            "(():,,)",
            "[][[]('comment is the gotcha':0.2,,);",
            "#SampleID\tHeaderA\tHeaderB\n0\t'yellow'\t0.45;",
            "))();",
            "((,,),((,,));"
        ]

    def test_newick_to_tree_node_valid_files(self):
        for tree, newicks in self.trees_newick_lists:
            for newick in newicks:
                fh = StringIO(newick)
                read_tree = _newick_to_tree_node(fh)

                self.assertTrue(self._is_equal(tree, read_tree))

                fh.close()

    def test_newick_to_tree_node_invalid_files(self):
        for invalid in self.invalid_newicks:
            fh = StringIO(invalid)
            with self.assertRaises(NewickFormatError):
                _newick_to_tree_node(fh)
            fh.close()

    def test_tree_node_to_newick(self):
        for tree, newicks in self.trees_newick_lists:
            newick = newicks[0]
            fh = StringIO()
            _tree_node_to_newick(tree, fh)

            self.assertEqual(newick, fh.getvalue())

            fh.close()

    def test_roundtrip(self):
        for tree, newicks in self.trees_newick_lists:
            newick = newicks[0]
            fh = StringIO(newick)
            tree = _newick_to_tree_node(fh)
            fh2 = StringIO()
            _tree_node_to_newick(tree, fh2)
            fh2.seek(0)
            tree2 = _newick_to_tree_node(fh2)

            self.assertEqual(newick, fh2.getvalue())
            self.assertTrue(self._is_equal(tree, tree2))

            fh.close()
            fh2.close()

    def test_newick_to_tree_node_convert_underscores(self):
        fh = StringIO('(_:0.1, _a, _b)__;')
        tree = _newick_to_tree_node(fh, convert_underscores=False)
        fh2 = StringIO()
        _tree_node_to_newick(tree, fh2)
        self.assertEquals(fh2.getvalue(), "('_':0.1,'_a','_b')'__';")
        fh2.close()
        fh.close()

    def test_newick_sniffer_valid_files(self):
        for _, newicks in self.trees_newick_lists:
            for newick in newicks:
                fh = StringIO(newick)
                self.assertEqual(_newick_sniffer(fh), (True, {}))
                fh.close()

    def test_newick_sniffer_invalid_files(self):
        for invalid in self.invalid_newicks:
            fh = StringIO(invalid)
            self.assertEqual(_newick_sniffer(fh), (False, {}))
            fh.close()

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