<commit_before>#!/usr/bin/env python3.3
# -*- coding: utf8 -*-
#
# Terminal component of CVE-Scan. Takes a list of scanned hosts as
# input and uses the terminal to displays this information in a
# graphical manner, enhancing it with information from the CVE-Search
# database and ToolsWatch Default Password Enumeration (DPE) list.
#
# Copyright (c) 2015	NorthernSec
# Copyright (c) 2015	Pieter-Jan Moreels
# This software is licensed under the Original BSD Licence

# Imports
import os
import sys
runpath = os.path.dirname(os.path.realpath(__file__))

import curses
from time import sleep
from lib.Toolkit import make_dict, enhance

# Class
class TermDisplay():
  @classmethod
  def start(self, systems=None):
    try:
      scr=curses.initscr()
      curses.start_color()
      curses.init_pair(1, curses.COLOR_RED,curses.COLOR_BLACK)
      curses.init_pair(2, curses.COLOR_BLUE, curses.COLOR_BLACK)
      curses.init_pair(3, curses.COLOR_WHITE, curses.COLOR_BLACK)
      curses.init_pair(4, curses.COLOR_GREEN, curses.COLOR_BLACK)
      curses.init_pair(5, curses.COLOR_YELLOW, curses.COLOR_BLACK)
      md_b   =curses.A_BOLD | curses.color_pair(1)
      md_bt  =curses.A_BOLD | curses.color_pair(2)
      md_cont=curses.A_BOLD | curses.A_BLINK
      md_head=curses.A_BOLD | curses.color_pair(2)
      md_norm=curses.color_pair(3)
      md_mark=curses.color_pair(1)
      md_bad=curses.color_pair(1)
      md_med=curses.color_pair(5)
      md_good=curses.color_pair(4)
      curses.curs_set(0)
      curses.noecho()
      scr.keypad(1)
    except Exception as ex:
      curses.endwin()
      raise Exception(ex)

    # Init sub functions
    def printWelcome():
      scr.clear()
      scr.border(1)
      maxy,maxx=scr.getmaxyx()
      if maxy<10:raise("Please make sure your terminal has at least 10 rows")
      anykey=int(((maxy-7)/2+8))
      start=int((maxx/2)-24)
      scr.addstr(2,start," _____ _   _ _____      _____                 ",md_b)
      scr.addstr(3,start,"/  __ \ | | |  ___|    /  ___|                ",md_b)
      scr.addstr(4,start,"| /  \/ | | | |__ _____\ `--.  ___ __ _ _ __  ",md_b)
      scr.addstr(5,start,"| |   | | | |  __|_____|`--. \/ __/ _` | '_ \ ",md_b)
      scr.addstr(6,start,"| \__/\ \_/ / |___     /\__/ / (_| (_| | | | |",md_b)
      scr.addstr(7,start," \____/\___/\____/     \____/ \___\__,_|_| |_|",md_b)
      scr.addstr(8,start,"                            (c) NorthernSec   ",md_bt)
      scr.addstr(anykey,start,"             [Press the any key]              ",md_cont)
      scr.refresh()
      scr.getch()

    def printHelp(section=None):
      if not section:
        text=[{'t':'----------','m':md_head},
              {'t':'|  HELP  |','m':md_head},
              {'t':'----------','m':md_head},
              {'t':' '},
              {'t':'Navigating','m':md_head},
              {'t':' * You can navigate through the scanned systems with p and n, or the left and right arrow key'},
              {'t':' * You can scroll through the current system with u and d, or the up and down arrow key'},
              {'t':' * You can jump directly to a scanned system by entering the page number with j, and entering the page number.'},
              {'t':' '},
              {'t':'Commands','m':md_head},
              {'t':'By pressing c, you can enter commands:'},
              {'t':'  help     Displays this menu'}]
      f=[' - press q to return to the previous page -']
      scroll(text,footer=f)
    def scroll(text,footer=[],cursor=False,limit=True):
      contInd=0
      cursInd=0
      cursPos=0
      while True:
        maxy,maxx=scr.getmaxyx()
        maxCont=maxy-len(footer)-4
        scr.clear()
        lines=list(text)
        if len(text)>maxCont:
          lines=lines[contInd:contInd+maxCont]

        for i,line in enumerate(lines):
          markup=line['m'] if 'm' in line else md_norm
          if 't' in line:
            scr.addstr(i+2,2,line['t'],markup)
          elif 'f' in line and 'v' in line:
            scr.addstr(i+2,2,line['f'],md_head)
            scr.addstr(i+2,14,str(line['v'])[:maxx-4],markup)
          else:
            scr.addstr(i+2,2,"Line is an invalid format")
        if cursor:
          scr.addstr(cursPos+2,1,">")
        for i,line in enumerate(footer):
          scr.addstr((maxy-len(footer)-1)+i,2,line)

        x=scr.getch()
        if x==ord('u') or x==curses.KEY_UP:
          if cursor:
            if cursInd>0:cursInd-=1
            if cursPos>0:cursPos-=1
            if cursPos==0 and contInd>0:contInd-=1
          else:
            if contInd>0:contInd-=1
        elif x==ord('d') or x==curses.KEY_DOWN:
          if cursor:
            if cursInd<len(text)-1:cursInd+=1
            if cursPos<maxCont and cursPos<len(text)-1:cursPos+=1
            if cursPos==maxCont and (len(text)-contInd)>maxCont:contInd+=1
          else:
            if (len(text)-contInd)>maxCont:contInd+=1
        elif x==ord('q'):
          return ord('q') if limit else (ord('q'),cursInd)
        else:
          if not limit:
            return (x,cursInd)

    def cvesForcpe(cpe,service=None):
      scroll(cpe)

    def product(banner):
      if banner:
        r=make_dict(banner)
        return r['product']
      else:
        return 'Unknown'

    def getLinesFor(sys):
      cpes=sys['cpes'] if 'cpes' in sys else ['Not Detected']
      mac=sys['mac'] if sys['mac'] else 'Unknown'
      marked=md_mark if 'cves' in cpes[0] else md_norm
      cont=[({'f':'IP',     'v':sys['ip']}),
            ({'f':'MAC',    'v':mac}),
            ({'f':'Status', 'v':sys['status']}),
            ({'f':'CPEs',   'v':cpes[0],       'm':marked})]
      if len(cpes)>1:
        for cpe in cpes[1:]:
          marked=md_mark if 'cves' in cpe else md_norm
          cont.append({'f':' ', 'v':cpe, 'm':marked})
      cont.append({'f':'Vendor', 'v':sys['vendor']})
      hosts=sys['hostnames'] if 'hostnames' in sys else ['None']
      cont.append({'f':'Hostnames', 'v':hosts[0]})
      for host in hosts[1:]:
        cont.append({'f':' ',  'v':host})
      cont.append({'f':'Distance', 'v':sys['distance']})
      serv=sys['services'] if 'services' in sys else ['No services found']
      s=serv[0]
      ser='%s (%s/%s) is %s'%(s['name'],s['port'],s['protocol'],s['state'])
      marked=md_mark if len(s['cves'])>0 else md_norm
      cont.append({'f':'Services', 'v':ser, 'm':marked})
      cont.append({'f':' ', 'v':' > %s'%product(s['banner']), 'm':marked})
      for s in serv[1:]:
        marked=md_mark if len(s['cves'])>0 else md_norm
        ser='%s (%s/%s) is %s'%(s['name'],s['port'],s['protocol'],s['state'])
        cont.append({'f':' ', 'v':ser, 'm':marked})
        cont.append({'f':' ', 'v':' > %s'%product(s['banner']), 'm':marked})

      return cont

    def cveView(cves):
      scroll(cves)

    def home():
      ind=0
      contInd=0
      while True:
        sys=systems[ind]
        content=getLinesFor(sys)
        foot=["(u)p   | (n)ext    | (p)revious | (q)uit",
              "(d)own | (j)ump to | (c)ommand  | [%s/%s]"%(ind+1,len(systems)),
              " "]
        x,index = scroll(content,footer=foot,cursor=True,limit=False)
        if x==ord('n') or x==curses.KEY_RIGHT:
          ind+=1
          if ind>=len(systems):ind=0
        elif x==ord('p') or x==curses.KEY_LEFT:
          ind-=1
          if ind<0:ind=len(systems)-1
        elif x==ord('q'):
          break
        elif x==ord('c'):
          maxy,maxx=scr.getmaxyx()
          curses.echo()
          scr.addstr(maxy-2,3,'>')
          line=scr.getstr(maxy-2,5).strip().lower()
          line=str(line,'utf-8')
          l=line.split(' ')
          c=l[0]
          args=l[1:]
          if c in ['help', 'h']:
            printHelp()
          elif c in ['cpe', 'c']:
            scr.addstr(2,2,content[index]['v'])
            scr.getch()
          elif c=='test':
            scr.addstr(2,2,str(args))
            scr.getch()
          else:
            error='Invalid Command'
            scr.addstr(int(maxy/2),int(maxx/2-len(error)/2),error,md_bad)
            scr.getch()
          curses.noecho()
        else:
          scr.addstr(2,2,("%s %s"%(x, chr(x))))
          scr.getch()

    try:
      printWelcome()
      home()
      curses.endwin()
    except Exception as ex:
      curses.endwin()
      raise Exception(ex)
<commit_msg>pattern recognision for cve browser functionality<commit_after>#!/usr/bin/env python3.3
# -*- coding: utf8 -*-
#
# Terminal component of CVE-Scan. Takes a list of scanned hosts as
# input and uses the terminal to displays this information in a
# graphical manner, enhancing it with information from the CVE-Search
# database and ToolsWatch Default Password Enumeration (DPE) list.
#
# Copyright (c) 2015	NorthernSec
# Copyright (c) 2015	Pieter-Jan Moreels
# This software is licensed under the Original BSD Licence

# Imports
import os
import sys
runpath = os.path.dirname(os.path.realpath(__file__))

import re
import curses
from time import sleep
from lib.Toolkit import make_dict, enhance

# Class
class TermDisplay():
  @classmethod
  def start(self, systems=None):
    try:
      scr=curses.initscr()
      curses.start_color()
      curses.init_pair(1, curses.COLOR_RED,curses.COLOR_BLACK)
      curses.init_pair(2, curses.COLOR_BLUE, curses.COLOR_BLACK)
      curses.init_pair(3, curses.COLOR_WHITE, curses.COLOR_BLACK)
      curses.init_pair(4, curses.COLOR_GREEN, curses.COLOR_BLACK)
      curses.init_pair(5, curses.COLOR_YELLOW, curses.COLOR_BLACK)
      md_b   =curses.A_BOLD | curses.color_pair(1)
      md_bt  =curses.A_BOLD | curses.color_pair(2)
      md_cont=curses.A_BOLD | curses.A_BLINK
      md_head=curses.A_BOLD | curses.color_pair(2)
      md_norm=curses.color_pair(3)
      md_mark=curses.color_pair(1)
      md_bad=curses.color_pair(1)
      md_med=curses.color_pair(5)
      md_good=curses.color_pair(4)
      curses.curs_set(0)
      curses.noecho()
      scr.keypad(1)
    except Exception as ex:
      curses.endwin()
      raise Exception(ex)

    # Init sub functions
    def printWelcome():
      scr.clear()
      scr.border(1)
      maxy,maxx=scr.getmaxyx()
      if maxy<10:raise("Please make sure your terminal has at least 10 rows")
      anykey=int(((maxy-7)/2+8))
      start=int((maxx/2)-24)
      scr.addstr(2,start," _____ _   _ _____      _____                 ",md_b)
      scr.addstr(3,start,"/  __ \ | | |  ___|    /  ___|                ",md_b)
      scr.addstr(4,start,"| /  \/ | | | |__ _____\ `--.  ___ __ _ _ __  ",md_b)
      scr.addstr(5,start,"| |   | | | |  __|_____|`--. \/ __/ _` | '_ \ ",md_b)
      scr.addstr(6,start,"| \__/\ \_/ / |___     /\__/ / (_| (_| | | | |",md_b)
      scr.addstr(7,start," \____/\___/\____/     \____/ \___\__,_|_| |_|",md_b)
      scr.addstr(8,start,"                            (c) NorthernSec   ",md_bt)
      scr.addstr(anykey,start,"             [Press the any key]              ",md_cont)
      scr.refresh()
      scr.getch()

    def printHelp(section=None):
      if not section:
        text=[{'t':'----------','m':md_head},
              {'t':'|  HELP  |','m':md_head},
              {'t':'----------','m':md_head},
              {'t':' '},
              {'t':'Navigating','m':md_head},
              {'t':' * You can navigate through the scanned systems with p and n, or the left and right arrow key'},
              {'t':' * You can scroll through the current system with u and d, or the up and down arrow key'},
              {'t':' * You can jump directly to a scanned system by entering the page number with j, and entering the page number.'},
              {'t':' '},
              {'t':'Commands','m':md_head},
              {'t':'By pressing c, you can enter commands:'},
              {'t':'  help     Displays this menu'}]
      f=[' - press q to return to the previous page -']
      scroll(text,footer=f)
    def scroll(text,footer=[],cursor=False,limit=True):
      contInd=0
      cursInd=0
      cursPos=0
      while True:
        maxy,maxx=scr.getmaxyx()
        maxCont=maxy-len(footer)-4
        scr.clear()
        lines=list(text)
        if len(text)>maxCont:
          lines=lines[contInd:contInd+maxCont]

        for i,line in enumerate(lines):
          markup=line['m'] if 'm' in line else md_norm
          if 't' in line:
            scr.addstr(i+2,2,line['t'],markup)
          elif 'f' in line and 'v' in line:
            scr.addstr(i+2,2,line['f'],md_head)
            scr.addstr(i+2,14,str(line['v'])[:maxx-4],markup)
          else:
            scr.addstr(i+2,2,"Line is an invalid format")
        if cursor:
          scr.addstr(cursPos+2,1,">")
        for i,line in enumerate(footer):
          scr.addstr((maxy-len(footer)-1)+i,2,line)

        x=scr.getch()
        if x==ord('u') or x==curses.KEY_UP:
          if cursor:
            if cursInd>0:cursInd-=1
            if cursPos>0:cursPos-=1
            if cursPos==0 and contInd>0:contInd-=1
          else:
            if contInd>0:contInd-=1
        elif x==ord('d') or x==curses.KEY_DOWN:
          if cursor:
            if cursInd<len(text)-1:cursInd+=1
            if cursPos<maxCont and cursPos<len(text)-1:cursPos+=1
            if cursPos==maxCont and (len(text)-contInd)>maxCont:contInd+=1
          else:
            if (len(text)-contInd)>maxCont:contInd+=1
        elif x==ord('q'):
          return ord('q') if limit else (ord('q'),cursInd)
        else:
          if not limit:
            return (x,cursInd)

    def cvesForcpe(cpe,service=None):
      scroll(cpe)

    def product(banner):
      if banner:
        r=make_dict(banner)
        return r['product']
      else:
        return 'Unknown'

    def getLinesFor(sys):
      cpes=sys['cpes'] if 'cpes' in sys else ['Not Detected']
      mac=sys['mac'] if sys['mac'] else 'Unknown'
      marked=md_mark if 'cves' in cpes[0] else md_norm
      cont=[({'f':'IP',     'v':sys['ip']}),
            ({'f':'MAC',    'v':mac}),
            ({'f':'Status', 'v':sys['status']}),
            ({'f':'CPEs',   'v':cpes[0],       'm':marked})]
      if len(cpes)>1:
        for cpe in cpes[1:]:
          marked=md_mark if 'cves' in cpe else md_norm
          cont.append({'f':' ', 'v':cpe, 'm':marked})
      cont.append({'f':'Vendor', 'v':sys['vendor']})
      hosts=sys['hostnames'] if 'hostnames' in sys else ['None']
      cont.append({'f':'Hostnames', 'v':hosts[0]})
      for host in hosts[1:]:
        cont.append({'f':' ',  'v':host})
      cont.append({'f':'Distance', 'v':sys['distance']})
      serv=sys['services'] if 'services' in sys else ['No services found']
      s=serv[0]
      ser='%s (%s/%s) is %s'%(s['name'],s['port'],s['protocol'],s['state'])
      marked=md_mark if len(s['cves'])>0 else md_norm
      cont.append({'f':'Services', 'v':ser, 'm':marked})
      cont.append({'f':' ', 'v':' > %s'%product(s['banner']), 'm':marked})
      cont.append({'f':' ', 'v':' > %s'%s['cpe'], 'm':marked})

      for s in serv[1:]:
        marked=md_mark if len(s['cves'])>0 else md_norm
        ser='%s (%s/%s) is %s'%(s['name'],s['port'],s['protocol'],s['state'])
        cont.append({'f':' ', 'v':ser, 'm':marked})
        cont.append({'f':' ', 'v':' > %s'%product(s['banner']), 'm':marked})
        cont.append({'f':' ', 'v':' > %s'%s['cpe'], 'm':marked})

      return cont

    def cveView(cves):
      scroll(cves)

    def home():
      ind=0
      contInd=0
      while True:
        sys=systems[ind]
        content=getLinesFor(sys)
        foot=["(u)p   | (n)ext    | (p)revious | (q)uit",
              "(d)own | (j)ump to | (c)ommand  | [%s/%s]"%(ind+1,len(systems)),
              " "]
        x,index = scroll(content,footer=foot,cursor=True,limit=False)
        if x==ord('n') or x==curses.KEY_RIGHT:
          ind+=1
          if ind>=len(systems):ind=0
        elif x==ord('p') or x==curses.KEY_LEFT:
          ind-=1
          if ind<0:ind=len(systems)-1
        elif x==ord('q'):
          break
        elif x==ord('c'):
          maxy,maxx=scr.getmaxyx()
          curses.echo()
          scr.addstr(maxy-2,3,'>')
          line=scr.getstr(maxy-2,5).strip().lower()
          line=str(line,'utf-8')
          l=line.split(' ')
          c=l[0]
          args=l[1:]
          if c in ['help', 'h']:
            printHelp()
          elif c in ['cpe', 'c']:
            # regexes to match
            rePort=re.compile('^([0-9]{1,4}|[1-5][0-9]{4}|6[0-4][0-9]{3}|65[0-4][0-9]{2}|655[0-2][0-9]|6553[0-5])$')
            reCPE=re.compile('^((%s|%s)[aoh]:.*)$'%(re.escape('cpe:/'),re.escape('cpe:2.3:')))
            if args:
              line=args[0]
            else:
              if content[index]['v'].startswith('> '):
                line=content[index]['v'].strip('> ')
              else:
                reg=re.compile('\\([0-9]*/(tcp|udp)\\)')
                line=reg.search(content[index]['v']).group()[1:-5]
            if rePort.match(line):
              # service = service where port equals
              scr.addstr(2,2,'port match')
            elif reCPE.match(line):
              # check if cpe matches service or os. if service, return service, if os, ?
              scr.addstr(2,2,'cpe match')
            else:
              # maybe a product. match services with product to find partial or full match
              scr.addstr(2,2,'no match')
            scr.getch()
          else:
            error='Invalid Command'
          if 'error' in locals() and error:
            scr.addstr(int(maxy/2),int(maxx/2-len(error)/2),error,md_bad)
            scr.getch()
            error=None
          curses.noecho()
        else:
          scr.addstr(2,2,("%s %s"%(x, chr(x))))
          scr.getch()

    try:
      printWelcome()
      home()
      curses.endwin()
    except Exception as ex:
      curses.endwin()
      raise Exception(ex)
<|endoftext|>


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

import sys
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.parser import Parser
from configparser import ConfigParser
from email.utils import parseaddr

#function definitions

def readlist(input_file):
    with open(input_file, mode="r") as f:
        lines = f.readlines()
    result = [i.strip() for i in lines]
    return result

def bounce(bouncetext, incoming):
    msg = MIMEMultipart()
    msg['Subject'] = "Re: " + incoming['subject']
    msg['From'] = incoming['to']
    msg['To'] = incoming['from']
    msg.attach(MIMEText(bouncetext, _charset='UTF-8'))
    s = smtplib.SMTP('localhost')
    s.send_message(msg)
    s.quit()
    exit(0)

#Parse email

email_in = sys.stdin.read()

incoming = Parser().parsestr(email_in)

sender = incoming['from']
sender_str = parseaddr(sender)[1]
to_list = parseaddr(incoming['to'])[1]
list_user = to_list.split('@')[0]
this_address = incoming['to']

#Parse config
basedir = '/scripts/mailforward/'

config = ConfigParser()
config.read(basedir + 'config.cfg', encoding="utf-8")

senders_file = basedir + config['FILES']['senders']
noreply_raw = config['RULES']['noreply'].split(',')
bounce_text = config['MESSAGES']['bounce']
attachment_text = config['MESSAGES']['attachment']
nolist_text = config['MESSAGES']['nolist']
noreply = []
for addr in noreply_raw:
    noreply.append(addr.strip())

#Try to read list addresses based on original recipient, bounce if no list found

try:
    list_file = basedir + list_user + ".list"
except FileNotFoundError:
    bounce(nolist_text)

#Try to read senders list, fallback to globals if not found
try:
    senders_file = basedir + list_user + "_senders.list"
except FileNotFoundError:
    senders_file = basedir + config['FILES']['globalsenders']

#Read senders and recipients into lists

senders = readlist(senders_file)
list_members = readlist(list_file)

if sender_str not in senders:

    if sender_str in noreply:
        exit(0)
    else:
        bounce(bounce_text, incoming)

else:

    for member in list_members:
        msg = MIMEMultipart()
        msg.set_payload(incoming)
        msg['From'] = this_address
        msg['reply-to'] = sender
        msg['To'] = member
        msg['Subject'] = incoming['subject']

        s = smtplib.SMTP('localhost')
        s.send_message(msg)
        s.quit()
<commit_msg>fix typo<commit_after>#!/usr/bin/env python3

import sys
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.parser import Parser
from configparser import ConfigParser
from email.utils import parseaddr

#function definitions

def readlist(input_file):
    with open(input_file, mode="r") as f:
        lines = f.readlines()
    result = [i.strip() for i in lines]
    return result

def bounce(bouncetext, incoming):
    msg = MIMEMultipart()
    msg['Subject'] = "Re: " + incoming['subject']
    msg['From'] = incoming['to']
    msg['To'] = incoming['from']
    msg.attach(MIMEText(bouncetext, _charset='UTF-8'))
    s = smtplib.SMTP('localhost')
    s.send_message(msg)
    s.quit()
    exit(0)

#Parse email

email_in = sys.stdin.read()

incoming = Parser().parsestr(email_in)

sender = incoming['from']
sender_str = parseaddr(sender)[1]
to_list = parseaddr(incoming['to'])[1]
list_user = to_list.split('@')[0]
this_address = incoming['to']

#Parse config
basedir = '/scripts/mailforward/'

config = ConfigParser()
config.read(basedir + 'config.cfg', encoding="utf-8")

noreply_raw = config['RULES']['noreply'].split(',')
bounce_text = config['MESSAGES']['bounce']
attachment_text = config['MESSAGES']['attachment']
nolist_text = config['MESSAGES']['nolist']
noreply = []
for addr in noreply_raw:
    noreply.append(addr.strip())

#Try to read list addresses based on original recipient, bounce if no list found

try:
    list_file = basedir + list_user + ".list"
except FileNotFoundError:
    bounce(nolist_text)

#Try to read senders list, fallback to globals if not found
try:
    senders_file = basedir + list_user + "_senders.list"
except FileNotFoundError:
    senders_file = basedir + config['FILES']['globalsenders']

#Read senders and recipients into lists

senders = readlist(senders_file)
list_members = readlist(list_file)

if sender_str not in senders:

    if sender_str in noreply:
        exit(0)
    else:
        bounce(bounce_text, incoming)

else:

    for member in list_members:
        msg = MIMEMultipart()
        msg.set_payload(incoming)
        msg['From'] = this_address
        msg['reply-to'] = sender
        msg['To'] = member
        msg['Subject'] = incoming['subject']

        s = smtplib.SMTP('localhost')
        s.send_message(msg)
        s.quit()
<|endoftext|>


<commit_before>import time
import sys
from display.display import Display
from fakesensor import FakeSensor
from sensors import Sensor
from display.graph import Graph
from display.clock import Clock
from teensy import Peripherals

class SleepApp:

	def __init__(self):
		self.quitting = False

		Peripherals.init()

		self.sensor = Sensor(Peripherals.devices["Pillow"])
		#self.sensor = FakeSensor()
		self.display = Display(320, 240)
		self.graph = Graph(self.display, self.sensor)
		self.clock = Clock(self.display)

	def start(self):
		lastTime = 0
		counter = 0
		startTime = time.time()
		while True:
			now = time.time()
			elapsed = now - lastTime
			if elapsed >= 0.005:
				lastTime = now
				self.sensor.readData()
				#self.clock.render()
				#self.graph.render((1,2,3))
				counter += 1
				if (counter % 100) == 0:
					print 100 / (time.time() - startTime)
					startTime = time.time()

	def quit(self):
		print "Quit..."
		self.quitting = True
		time.sleep(1)
		sys.exit()
<commit_msg>Peripherals<commit_after>import time
import sys
from display.display import Display
from fakesensor import FakeSensor
from sensors import Sensor
from display.graph import Graph
from display.clock import Clock
from teensy import Peripherals

class SleepApp:

	def __init__(self):
		self.quitting = False

		Peripherals.init()

		self.sensor = Sensor(Peripherals.devices["Pillow"])
		#self.sensor = FakeSensor()
		self.display = Display(320, 240)
		self.graph = Graph(self.display, self.sensor)
		self.clock = Clock(self.display)

	def start(self):
		lastTime = 0
		counter = 0
		startTime = time.time()
		while True:
			now = time.time()
			elapsed = now - lastTime
			if elapsed >= 0.001:
				lastTime = now
				self.sensor.readData()
				self.clock.render()
				self.graph.render((1,2,3))
				counter += 1
				if (counter % 100) == 0:
					print 100 / (time.time() - startTime)
					startTime = time.time()

	def quit(self):
		print "Quit..."
		self.quitting = True
		time.sleep(1)
		sys.exit()
<|endoftext|>


<commit_before>import h2o, h2o_cmd, re, os

# hdfs/maprfs/s3/s3n paths should be absolute from the bucket (top level)
# so only walk around for local
def find_folder_and_filename(bucket, pathWithRegex):
    checkPath = True
    # strip the common mistake of leading "/" in path, if bucket is specified too
    if bucket is not None and re.match("/", pathWithRegex):
        h2o.verboseprint("You said bucket:", bucket, "so stripping incorrect leading '/' from", pathWithRegex)
        pathWithRegex = pathWithRegex.lstrip('/')

    if bucket is None:  # good for absolute path name
        bucketPath = ""

    elif bucket == ".":
        bucketPath = os.getcwd()

    # only use if the build_cloud was for remote H2O
    # Never use the var for remote, if you're doing a put! (which always sources local)
    elif h2o.nodes[0].remoteH2O and schema!='put' and os.environ.get('H2O_REMOTE_BUCKETS_ROOT'):
        # we may use this to force remote paths, so don't look locally for file
        rootPath = os.environ.get('H2O_REMOTE_BUCKETS_ROOT')
        bucketPath = os.path.join(rootPath, bucket)
        checkpath = False

    # does it work to use bucket "." to get current directory
    elif os.environ.get('H2O_BUCKETS_ROOT'):
        rootPath = os.environ.get('H2O_BUCKETS_ROOT')
        print "Using H2O_BUCKETS_ROOT environment variable:", rootPath

        if not (os.path.exists(rootPath)):
            raise Exception("H2O_BUCKETS_ROOT in env but %s doesn't exist." % rootPath)

        bucketPath = os.path.join(rootPath, bucket)
        if not (os.path.exists(bucketPath)):
            raise Exception("H2O_BUCKETS_ROOT and path used to form %s which doesn't exist." % bucketPath)

    else:
        # if we run remotely, we're assuming the import folder path on the remote machine
        # matches what we find on our local machine. But maybe the local user doesn't exist remotely
        # so using his path won't work.
        # Resolve by looking for special state in the config. If user = 0xdiag, just force the bucket location
        # This is a lot like knowing about fixed paths with s3 and hdfs
        # Otherwise the remote path needs to match the local discovered path.

        # want to check the username being used remotely first. should exist here too if going to use
        possibleUsers = ["~"]
        h2o.verboseprint("h2o.nodes[0].username:", h2o.nodes[0].username)
        if h2o.nodes[0].username:
            possibleUsers.insert(0, "~" + h2o.nodes[0].username)

        for u in possibleUsers:
            rootPath = os.path.expanduser(u)
            bucketPath = os.path.join(rootPath, bucket)
            h2o.verboseprint("Checking bucketPath:", bucketPath, 'assuming home is', rootPath)
            if os.path.exists(bucketPath):
                h2o.verboseprint("search A did find", bucket, "at", rootPath)
                break
        else:
            # last chance to find it by snooping around
            rootPath = os.getcwd()
            h2o.verboseprint("find_bucket looking upwards from", rootPath, "for", bucket)
            # don't spin forever
            levels = 0
            while not (os.path.exists(os.path.join(rootPath, bucket))):
                h2o.verboseprint("Didn't find", bucket, "at", rootPath)
                rootPath = os.path.split(rootPath)[0]
                levels += 1
                if (levels==6):
                    raise Exception("unable to find bucket: %s" % bucket)

            h2o.verboseprint("search B did find", bucket, "at", rootPath)
            bucketPath = os.path.join(rootPath, bucket)

    # if there's no path, just return the bucketPath
    # but what about cases with a header in the folder too? (not putfile)
    if pathWithRegex is None:
        return (bucketPath, None)

    # if there is a "/" in the path, that means it's not just a pattern
    # split it
    # otherwise it is a pattern. use it to search for files in python first?
    # FIX! do that later
    elif "/" in pathWithRegex:
        (head, tail) = os.path.split(pathWithRegex)
        folderPath = os.path.abspath(os.path.join(bucketPath, head))
        if checkPath and not os.path.exists(folderPath):
            raise Exception("%s doesn't exist. %s under %s may be wrong?" % (folderPath, head, bucketPath))
    else:
        folderPath = bucketPath
        tail = pathWithRegex

    h2o.verboseprint("folderPath:", folderPath, "tail:", tail)
    return (folderPath, tail)

# passes additional params thru kwargs for parse
# use_header_file=
# header=
# exclude=
# src_key= only used if for put file key name (optional)
# path should point to a file or regex of files. (maybe folder works? but unnecessary
def import_only(node=None, schema='local', bucket=None, path=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, doSummary=True, src_key='python_src_key', **kwargs):

    # no bucket is sometimes legal (fixed path)
    if not node: node = h2o.nodes[0]

    if "/" in path:
        (head, pattern) = os.path.split(path)
    else:
        (head, pattern)  = ("", path)

    h2o.verboseprint("head:", head)
    h2o.verboseprint("pattern:", pattern)

    # to train users / okay here
    if re.search(r"[\*<>{}[\]~`]", head):
       raise Exception("h2o folder path %s can't be regex. path= was %s" % (head, path))

    if schema=='put':
        # to train users
        if re.search(r"[/\*<>{}[\]~`]", pattern):
           raise Exception("h2o putfile basename %s can't be regex. path= was %s" % (pattern, path))

        if not path:
            raise Exception("path= didn't say what file to put")

        (folderPath, filename) = find_folder_and_filename(bucket, path)
        h2o.verboseprint("folderPath:", folderPath, "filename:", filename)
        filePath = os.path.join(folderPath, filename)
        h2o.verboseprint('filePath:', filePath)
        key = node.put_file(filePath, key=src_key, timeoutSecs=timeoutSecs)
        return (None, key)

    if schema=='local' and not \
            (node.redirect_import_folder_to_s3_path or node.redirect_import_folder_to_s3n_path):
        (folderPath, pattern) = find_folder_and_filename(bucket, path)
        folderURI = 'nfs:/' + folderPath
        importResult = node.import_files(folderPath, timeoutSecs=timeoutSecs)

    else:
        if bucket is not None and re.match("/", head):
            h2o.verboseprint("You said bucket:", bucket, "so stripping incorrect leading '/' from", head)
            head = head.lstrip('/')

        # strip leading / in head if present
        if bucket and head!="":
            folderOffset = bucket + "/" + head
        elif bucket:
            folderOffset = bucket
        else:
            folderOffset = head

        if schema=='s3' or node.redirect_import_folder_to_s3_path:
            folderURI = "s3://" + folderOffset
            importResult = node.import_s3(bucket, timeoutSecs=timeoutSecs)

        elif schema=='s3n' or node.redirect_import_folder_to_s3n_path:
            folderURI = "s3n://" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        elif schema=='maprfs':
            folderURI = "hdfs:///" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        elif schema=='hdfs':
            h2o.verboseprint(h2o.nodes[0].hdfs_name_node)
            h2o.verboseprint("folderOffset;", folderOffset)
            folderURI = "hdfs://" + h2o.nodes[0].hdfs_name_node + "/" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        else:
            raise Exception("schema not understood: %s" % schema)

    importPattern = folderURI + "/" + pattern
    return (importResult, importPattern)

# can take header, header_from_file, exclude params
def parse_only(node=None, pattern=None, hex_key=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, **kwargs):

    if not node: node = h2o.nodes[0]

    parseResult = node.parse(key=pattern, key2=hex_key,
        timeoutSecs=timeoutSecs, retryDelaySecs=retryDelaySecs,
        initialDelaySecs=initialDelaySecs, pollTimeoutSecs=pollTimeoutSecs, noise=noise,
        benchmarkLogging=None, noPoll=noPoll, **kwargs)

    parseResult['python_source'] = pattern
    return parseResult

def import_parse(node=None, schema='local', bucket=None, path=None,
    src_key=None, hex_key=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, doSummary=True, **kwargs):

    if not node: node = h2o.nodes[0]

    (importResult, importPattern) = import_only(node, schema, bucket, path,
        timeoutSecs, retryDelaySecs, initialDelaySecs, pollTimeoutSecs, noise,
        benchmarkLogging, noPoll, doSummary, src_key, **kwargs)

    h2o.verboseprint("importPattern:", importPattern)
    h2o.verboseprint("importResult", h2o.dump_json(importResult))

    parseResult = parse_only(node, importPattern, hex_key,
        timeoutSecs, retryDelaySecs, initialDelaySecs, pollTimeoutSecs, noise,
        benchmarkLogging, noPoll, **kwargs)
    h2o.verboseprint("parseResult:", h2o.dump_json(parseResult))

    # do SummaryPage here too, just to get some coverage
    if doSummary:
        node.summary_page(parseResult['destination_key'], timeoutSecs=timeoutSecs)

    return parseResult

# returns full key name, from current store view
def find_key(pattern=None):
    found = None
    kwargs = {'filter': None}
    storeViewResult = h2o.nodes[0].store_view(**kwargs)
    keys = storeViewResult['keys']
    if len(keys) == 0:
        return None

    if len(keys) > 1:
        h2o.verboseprint("Warning: multiple imported keys match the key pattern given, Using: %s" % keys[0]['key'])

    return keys[0]['key']

# the storeViewResult for every node may or may not be the same
# supposed to be the same? In any case
# pattern can't be regex to h2o?
# None should be same as no pattern
def delete_keys(node=None, pattern=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    kwargs = {'filter': pattern}
    storeViewResult = h2o_cmd.runStoreView(node, timeoutSecs=timeoutSecs, **kwargs)
    keys = storeViewResult['keys']
    for k in keys:
        node.remove_key(k['key'])
    deletedCnt = len(keys)
    print "Deleted", deletedCnt, "keys at", node
    return deletedCnt

def delete_keys_at_all_nodes(node=None, pattern=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    totalDeletedCnt = 0
    # do it in reverse order, since we always talk to 0 for other stuff
    # this will be interesting if the others don't have a complete set
    # theoretically, the deletes should be 0 after the first node
    # since the deletes should be global
    for node in reversed(h2o.nodes):
        deletedCnt = delete_keys(node, pattern=pattern, timeoutSecs=timeoutSecs)
        totalDeletedCnt += deletedCnt
    print "\nTotal: Deleted", totalDeletedCnt, "keys at", len(h2o.nodes), "nodes"
    return totalDeletedCnt

# Since we can't trust a single node storeview list, this will get keys that match text
# for deleting, from a list saved from an import
def delete_keys_from_import_result(node=None, pattern=None, importResult=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    # the list could be from hdfs/s3 or local. They have to different list structures
    deletedCnt = 0
    if 'succeeded' in importResult:
        kDict = importResult['succeeded']
        for k in kDict:
            key = k['key']
            if (pattern in key) or pattern is None:
                print "\nRemoving", key
                removeKeyResult = node.remove_key(key=key)
                deletedCnt += 1
    elif 'keys' in importResult:
        kDict = importResult['keys']
        for k in kDict:
            key = k
            if (pattern in key) or pattern is None:
                print "\nRemoving", key
                removeKeyResult = node.remove_key(key=key)
                deletedCnt += 1
    else:
        raise Exception ("Can't find 'files' or 'succeeded' in your file dict. why? not from hdfs/s3 or local?")
    print "Deleted", deletedCnt, "keys at", node
    return deletedCnt
<commit_msg>can also find buckets or symlinks to buckets at home dirs of username, h2o username, or 0xdiag<commit_after>import h2o, h2o_cmd, re, os
import getpass

# hdfs/maprfs/s3/s3n paths should be absolute from the bucket (top level)
# so only walk around for local
def find_folder_and_filename(bucket, pathWithRegex):
    checkPath = True
    # strip the common mistake of leading "/" in path, if bucket is specified too
    if bucket is not None and re.match("/", pathWithRegex):
        h2o.verboseprint("You said bucket:", bucket, "so stripping incorrect leading '/' from", pathWithRegex)
        pathWithRegex = pathWithRegex.lstrip('/')

    if bucket is None:  # good for absolute path name
        bucketPath = ""

    elif bucket == ".":
        bucketPath = os.getcwd()

    # only use if the build_cloud was for remote H2O
    # Never use the var for remote, if you're doing a put! (which always sources local)
    elif h2o.nodes[0].remoteH2O and schema!='put' and os.environ.get('H2O_REMOTE_BUCKETS_ROOT'):
        # we may use this to force remote paths, so don't look locally for file
        rootPath = os.environ.get('H2O_REMOTE_BUCKETS_ROOT')
        bucketPath = os.path.join(rootPath, bucket)
        checkpath = False

    # does it work to use bucket "." to get current directory
    elif os.environ.get('H2O_BUCKETS_ROOT'):
        rootPath = os.environ.get('H2O_BUCKETS_ROOT')
        print "Using H2O_BUCKETS_ROOT environment variable:", rootPath

        if not (os.path.exists(rootPath)):
            raise Exception("H2O_BUCKETS_ROOT in env but %s doesn't exist." % rootPath)

        bucketPath = os.path.join(rootPath, bucket)
        if not (os.path.exists(bucketPath)):
            raise Exception("H2O_BUCKETS_ROOT and path used to form %s which doesn't exist." % bucketPath)

    else:
        # if we run remotely, we're assuming the import folder path on the remote machine
        # matches what we find on our local machine. But maybe the local user doesn't exist remotely
        # so using his path won't work.
        # Resolve by looking for special state in the config. If user = 0xdiag, just force the bucket location
        # This is a lot like knowing about fixed paths with s3 and hdfs
        # Otherwise the remote path needs to match the local discovered path.

        # want to check the username being used remotely first. should exist here too if going to use
        username = getpass.getuser()
        h2oUsername = h2o.nodes[0].username
        h2o.verboseprint("username:", username, "h2oUsername:", h2oUsername)
        # resolved in order, looking for bucket (ln -s will work) in these home dirs.
        if h2oUsername != username:
            possibleUsers = [username, h2oUsername, "0xdiag"]
        else:
            possibleUsers = [username, "0xdiag"]

        for u in possibleUsers:
            rootPath = os.path.expanduser("~" + u)
            bucketPath = os.path.join(rootPath, bucket)
            h2o.verboseprint("Checking bucketPath:", bucketPath, 'assuming home is', rootPath)
            if os.path.exists(bucketPath):
                h2o.verboseprint("search A did find", bucket, "at", rootPath)
                break
        else:
            # last chance to find it by snooping around
            rootPath = os.getcwd()
            h2o.verboseprint("find_bucket looking upwards from", rootPath, "for", bucket)
            # don't spin forever
            levels = 0
            while not (os.path.exists(os.path.join(rootPath, bucket))):
                h2o.verboseprint("Didn't find", bucket, "at", rootPath)
                rootPath = os.path.split(rootPath)[0]
                levels += 1
                if (levels==6):
                    raise Exception("unable to find bucket: %s" % bucket)

            h2o.verboseprint("search B did find", bucket, "at", rootPath)
            bucketPath = os.path.join(rootPath, bucket)

    # if there's no path, just return the bucketPath
    # but what about cases with a header in the folder too? (not putfile)
    if pathWithRegex is None:
        return (bucketPath, None)

    # if there is a "/" in the path, that means it's not just a pattern
    # split it
    # otherwise it is a pattern. use it to search for files in python first?
    # FIX! do that later
    elif "/" in pathWithRegex:
        (head, tail) = os.path.split(pathWithRegex)
        folderPath = os.path.abspath(os.path.join(bucketPath, head))
        if checkPath and not os.path.exists(folderPath):
            raise Exception("%s doesn't exist. %s under %s may be wrong?" % (folderPath, head, bucketPath))
    else:
        folderPath = bucketPath
        tail = pathWithRegex

    h2o.verboseprint("folderPath:", folderPath, "tail:", tail)
    return (folderPath, tail)

# passes additional params thru kwargs for parse
# use_header_file=
# header=
# exclude=
# src_key= only used if for put file key name (optional)
# path should point to a file or regex of files. (maybe folder works? but unnecessary
def import_only(node=None, schema='local', bucket=None, path=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, doSummary=True, src_key='python_src_key', **kwargs):

    # no bucket is sometimes legal (fixed path)
    if not node: node = h2o.nodes[0]

    if "/" in path:
        (head, pattern) = os.path.split(path)
    else:
        (head, pattern)  = ("", path)

    h2o.verboseprint("head:", head)
    h2o.verboseprint("pattern:", pattern)

    # to train users / okay here
    if re.search(r"[\*<>{}[\]~`]", head):
       raise Exception("h2o folder path %s can't be regex. path= was %s" % (head, path))

    if schema=='put':
        # to train users
        if re.search(r"[/\*<>{}[\]~`]", pattern):
           raise Exception("h2o putfile basename %s can't be regex. path= was %s" % (pattern, path))

        if not path:
            raise Exception("path= didn't say what file to put")

        (folderPath, filename) = find_folder_and_filename(bucket, path)
        h2o.verboseprint("folderPath:", folderPath, "filename:", filename)
        filePath = os.path.join(folderPath, filename)
        h2o.verboseprint('filePath:', filePath)
        key = node.put_file(filePath, key=src_key, timeoutSecs=timeoutSecs)
        return (None, key)

    if schema=='local' and not \
            (node.redirect_import_folder_to_s3_path or node.redirect_import_folder_to_s3n_path):
        (folderPath, pattern) = find_folder_and_filename(bucket, path)
        folderURI = 'nfs:/' + folderPath
        importResult = node.import_files(folderPath, timeoutSecs=timeoutSecs)

    else:
        if bucket is not None and re.match("/", head):
            h2o.verboseprint("You said bucket:", bucket, "so stripping incorrect leading '/' from", head)
            head = head.lstrip('/')

        # strip leading / in head if present
        if bucket and head!="":
            folderOffset = bucket + "/" + head
        elif bucket:
            folderOffset = bucket
        else:
            folderOffset = head

        if schema=='s3' or node.redirect_import_folder_to_s3_path:
            folderURI = "s3://" + folderOffset
            importResult = node.import_s3(bucket, timeoutSecs=timeoutSecs)

        elif schema=='s3n' or node.redirect_import_folder_to_s3n_path:
            folderURI = "s3n://" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        elif schema=='maprfs':
            folderURI = "hdfs:///" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        elif schema=='hdfs':
            h2o.verboseprint(h2o.nodes[0].hdfs_name_node)
            h2o.verboseprint("folderOffset;", folderOffset)
            folderURI = "hdfs://" + h2o.nodes[0].hdfs_name_node + "/" + folderOffset
            importResult = node.import_hdfs(folderURI, timeoutSecs=timeoutSecs)

        else:
            raise Exception("schema not understood: %s" % schema)

    importPattern = folderURI + "/" + pattern
    return (importResult, importPattern)

# can take header, header_from_file, exclude params
def parse_only(node=None, pattern=None, hex_key=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, **kwargs):

    if not node: node = h2o.nodes[0]

    parseResult = node.parse(key=pattern, key2=hex_key,
        timeoutSecs=timeoutSecs, retryDelaySecs=retryDelaySecs,
        initialDelaySecs=initialDelaySecs, pollTimeoutSecs=pollTimeoutSecs, noise=noise,
        benchmarkLogging=None, noPoll=noPoll, **kwargs)

    parseResult['python_source'] = pattern
    return parseResult

def import_parse(node=None, schema='local', bucket=None, path=None,
    src_key=None, hex_key=None,
    timeoutSecs=30, retryDelaySecs=0.5, initialDelaySecs=0.5, pollTimeoutSecs=180, noise=None,
    benchmarkLogging=None, noPoll=False, doSummary=True, **kwargs):

    if not node: node = h2o.nodes[0]

    (importResult, importPattern) = import_only(node, schema, bucket, path,
        timeoutSecs, retryDelaySecs, initialDelaySecs, pollTimeoutSecs, noise,
        benchmarkLogging, noPoll, doSummary, src_key, **kwargs)

    h2o.verboseprint("importPattern:", importPattern)
    h2o.verboseprint("importResult", h2o.dump_json(importResult))

    parseResult = parse_only(node, importPattern, hex_key,
        timeoutSecs, retryDelaySecs, initialDelaySecs, pollTimeoutSecs, noise,
        benchmarkLogging, noPoll, **kwargs)
    h2o.verboseprint("parseResult:", h2o.dump_json(parseResult))

    # do SummaryPage here too, just to get some coverage
    if doSummary:
        node.summary_page(parseResult['destination_key'], timeoutSecs=timeoutSecs)

    return parseResult

# returns full key name, from current store view
def find_key(pattern=None):
    found = None
    kwargs = {'filter': None}
    storeViewResult = h2o.nodes[0].store_view(**kwargs)
    keys = storeViewResult['keys']
    if len(keys) == 0:
        return None

    if len(keys) > 1:
        h2o.verboseprint("Warning: multiple imported keys match the key pattern given, Using: %s" % keys[0]['key'])

    return keys[0]['key']

# the storeViewResult for every node may or may not be the same
# supposed to be the same? In any case
# pattern can't be regex to h2o?
# None should be same as no pattern
def delete_keys(node=None, pattern=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    kwargs = {'filter': pattern}
    storeViewResult = h2o_cmd.runStoreView(node, timeoutSecs=timeoutSecs, **kwargs)
    keys = storeViewResult['keys']
    for k in keys:
        node.remove_key(k['key'])
    deletedCnt = len(keys)
    print "Deleted", deletedCnt, "keys at", node
    return deletedCnt

def delete_keys_at_all_nodes(node=None, pattern=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    totalDeletedCnt = 0
    # do it in reverse order, since we always talk to 0 for other stuff
    # this will be interesting if the others don't have a complete set
    # theoretically, the deletes should be 0 after the first node
    # since the deletes should be global
    for node in reversed(h2o.nodes):
        deletedCnt = delete_keys(node, pattern=pattern, timeoutSecs=timeoutSecs)
        totalDeletedCnt += deletedCnt
    print "\nTotal: Deleted", totalDeletedCnt, "keys at", len(h2o.nodes), "nodes"
    return totalDeletedCnt

# Since we can't trust a single node storeview list, this will get keys that match text
# for deleting, from a list saved from an import
def delete_keys_from_import_result(node=None, pattern=None, importResult=None, timeoutSecs=30):
    if not node: node = h2o.nodes[0]
    # the list could be from hdfs/s3 or local. They have to different list structures
    deletedCnt = 0
    if 'succeeded' in importResult:
        kDict = importResult['succeeded']
        for k in kDict:
            key = k['key']
            if (pattern in key) or pattern is None:
                print "\nRemoving", key
                removeKeyResult = node.remove_key(key=key)
                deletedCnt += 1
    elif 'keys' in importResult:
        kDict = importResult['keys']
        for k in kDict:
            key = k
            if (pattern in key) or pattern is None:
                print "\nRemoving", key
                removeKeyResult = node.remove_key(key=key)
                deletedCnt += 1
    else:
        raise Exception ("Can't find 'files' or 'succeeded' in your file dict. why? not from hdfs/s3 or local?")
    print "Deleted", deletedCnt, "keys at", node
    return deletedCnt
<|endoftext|>


<commit_before>"""A module for defining structured supply species.

:author: Matthew Gidden <matthew.gidden _at_ gmail.com>
"""
import itertools
import numpy as np
import random
import math

from collections import OrderedDict, defaultdict, Iterable, namedtuple

from cyclopts import tools as cyctools
from cyclopts import cyclopts_io as cycio
import cyclopts.exchange_instance as exinst
from cyclopts.problems import ProblemSpecies
from cyclopts.exchange_family import ResourceExchange

from cyclopts.structured_species import data
from cyclopts.structured_species import tools as strtools
from cyclopts.structured_species import request

class Point(strtools.Point):
    """A container class representing a point in parameter space"""

    """ordered mapping from input parameters to default values and np.dtypes, see
    the theory manual for further explanation of the parameter names"""
    parameters = OrderedDict(sorted(request.Point.parameters.items() + {
                "d_th": strtools.Param([0.67, 0.33, 0], (np.float64, 3)),
                "d_f_mox": strtools.Param([0., 0., 1., 0.], (np.float64, 4)),
                "d_f_thox": strtools.Param([0., 0., 0., 1.], (np.float64, 4)),
                "f_repo": strtools.Param(0.1, np.float32),
                }.items(), key=lambda t: t[0]))

    def __init__(self, d=None):
        """Parameters
        ----------
        d : dict, optional
            a dictionary with key value pairs of parameter name, parameter
            value
        """
        super(Point, self).__init__(d)
        if self.seed > 0:
            random.seed(self.seed)

    def _parameters(self):
        """ordered mapping from input parameters to default values and np.dtypes, see
        the theory manual for further explanation of the parameter names"""
        return Point.parameters

class Reactor(strtools.Reactor):
    """An extension reactor model for Structured Supply Species"""
    def __init__(self, kind, point=None, n_assems=None):
        super(Reactor, self).__init__(kind, point, n_assems)
        self.assem_qty = data.fuel_unit * data.core_vol_frac[self.kind] \
            / self.n_assems

    def gen_group(self, gid):
        supply = False
        grp = exinst.ExGroup(gid, supply)
        grp.AddCap(self.assem_qty)
        return grp

    def gen_node(self, nid, gid, excl_id):
        supply = False
        excl = True
        return exinst.ExNode(nid, gid, supply,
                             self.assem_qty, excl, excl_id)

class Requester(object):
    """A simplified requester model for Structured Supply Species"""
    def __init__(self, kind, gids, nids):
        self.kind = kind
        self.req_qty = data.sup_rhs[self.kind]
        gid = gids.next()
        req = True
        if self.kind == data.Supports.repo:
            self.group = grp = exinst.ExGroup(gid, req, self.req_qty)
        else:
            self.group = grp = exinst.ExGroup(gid, req, self.req_qty)
            commod = data.sup_to_commod[self.kind]
            rxtr = data.sup_to_rxtr[self.kind]
            grp.AddCap(self.req_qty * strtools.mean_enr(rxtr, commod) / 100 \
                           * data.relative_qtys[rxtr][commod])
        self._gen_nodes(gid, nids)
        self.loc = data.loc()

    def _gen_nodes(self, gid, nids):
        self.nodes = []
        self.commod_to_nodes = {}
        req = True
        for commod in data.sup_pref_basis[self.kind].keys():
            nid = nids.next()
            node = exinst.ExNode(nid, gid, req, self.req_qty)
            self.nodes.append(node)
            self.commod_to_nodes[commod] = node

    def coeff(self, enr, rkind, commod):
        if self.kind == data.Supports.repo:
            raise RuntimeError('Coeff not supported for repos')
        return enr / 100 * data.relative_qtys[rkind][commod]

class StructuredSupply(ProblemSpecies):
    """A class representing structured supply-based exchanges species."""

    @staticmethod
    def pnt_to_realization(point):
        """Returns a realization of a structured supply instance given a point
        in parameter space.

        A realization is a namedtuple of :
          * reqrs: a dictionary of the kind and number of each requester
          * rxtrs: a dictionary of the kind and number of each reactor
          * assem_dists: a dictionary of the kind of reactor to a dictionary
            of Commodity type to the number of assemblies of that Commodity type
        """
        reqrs = {data.Supports[i]: n \
                     for i, n in enumerate(strtools.support_breakdown(point))}
        rxtrs = {data.Reactor[i]: n \
                     for i, n in enumerate(strtools.reactor_breakdown(point))}
        dists = {k: strtools.assembly_breakdown(point, k) \
                     for k in data.Reactors}
        keys = ['n_reqrs', 'n_rxtrs', 'assem_dists']
        return namedtuple('Realization', keys)(reqrs, rxtrs, dists)

    @staticmethod
    def gen_arc(aid, point, commod, rx_node_id, rxtr, reqr):
        """generate an arc"""
        pref = strtools.preference(data.sup_pref_basis[reqr.kind][commod],
                                   rxtr.loc, reqr.loc,
                                   point.f_loc, point.r_l_c, point.n_reg)
        rq_coeffs = [reqr.coeff(rxtr.enr(commod), rxtr.kind, commod)] \
            if not reqr.kind == data.Supports.repo else []
        arc = exinst.ExArc(aid,
                           reqr.commod_to_nodes[commod].id, rq_coeffs,
                           rx_node_id, [1],
                           pref)
        return arc

    def __init__(self):
        super(StructuredSupply, self).__init__()
        self._family = ResourceExchange()
        self.space = None
        self._n_points = None
        # 16 bytes for uuid
        self.param_tbl_name = 'Points'
        self._param_dtype = np.dtype(
            [('paramid', ('str', 16)), ('family', ('str', 30))] + \
                [(name, param.dtype) for name, param in Point.parameters.items()])
        self.sum_tbl_name = 'Summary'
        facs = ['n_r_th', 'n_r_f_mox', 'n_r_f_thox', 'n_s_uox', 'n_s_th_mox',
                'n_s_f_mox', 'n_s_f_thox', 'n_s_repo']
        self.iter_params = ['d_th', 'd_f_mox', 'd_f_thox']
        self._sum_dtype = np.dtype(
            [('paramid', ('str', 16)), ('family', ('str', 30))] + \
                [(name, np.uint32) for name in facs])
        # reset id generation
        self.nids = cyctools.Incrementer()
        self.excl_ids = cyctools.Incrementer()
        self.gids = cyctools.Incrementer()
        self.arcids = cyctools.Incrementer()

        # default realization is None
        self._rlztn = None

    @property
    def family(self):
        """Returns
        -------
        family : ResourceExchange
            An instance of this species' family
        """
        return self._family

    @property
    def name(self):
        """Returns
        -------
        name : string
            The name of this species
        """
        return 'StructuredSupply'

    def register_tables(self, h5file, prefix):
        """Parameters
        ----------
        h5file : PyTables File
            the hdf5 file
        prefix : string
            the absolute path to the group for tables of this species

        Returns
        -------
        tables : list of cyclopts_io.Tables
            All tables that could be written to by this species.
        """
        return [cycio.Table(h5file, '/'.join([prefix, self.param_tbl_name]),
                            self._param_dtype),
                cycio.Table(h5file, '/'.join([prefix, self.sum_tbl_name]),
                            self._sum_dtype)]

    def read_space(self, space_dict):
        """Parameters
        ----------
        space_dict : dict
            A dictionary container resulting from the reading in of a run
            control file
        """
        self.space = {k: v if isinstance(v, Iterable) else [v] \
                          for k, v in space_dict.items() \
                          if k in Point.parameters}

    @property
    def n_points(self):
        """Returns
        -------
        n : int
            The total number of points in the parameter space
        """
        return cyctools.n_permutations(self.space, iter_keys=self.iter_params)

    def points(self):
        """Derived classes must implement this function returning a
        representation of a point in its parameter space to be used by other
        class member functions.

        Returns
        -------
        point_generator : generator
            A generator for representation of a point in parameter space to be
            used by this species
        """
        keys = self.space.keys()
        vals = self.space.values()
        for args in cyctools.expand_args(vals):
            d = {keys[i]: args[i] for i in range(len(args))}
            yield Point(d)

    def record_point(self, point, param_uuid, tables):
        """Parameters
        ----------
        point : tuple or other
            A representation of a point in parameter space
        param_uuid : uuid
            The uuid of the point in parameter space
        tables : list of cyclopts_io.Table
            The tables that can be written to
        """
        uid = param_uuid.bytes if len(param_uuid.bytes) == 16 \
            else param_uuid.bytes + '\0'

        data = [param_uuid.bytes, self._family.name]
        data += [getattr(point, k) for k in Point.parameters.keys()]
        tables[self.param_tbl_name].append_data([tuple(data)])

        data = [param_uuid.bytes, self._family.name]
        data += strtools.reactor_breakdown(point)
        data += strtools.support_breakdown(point)
        tables[self.sum_tbl_name].append_data([tuple(data)])

    def _get_reactors(self):
        # requires self._rlztn to be set
        rkinds = self._rlztn.n_rxtrs.keys()
        n_assems = {k: sum(v.values()) \
                       for k, v in self._rlztn.assem_dists.items()}
        gen_ary = lambda kind, num, n_assems: \
            np.ndarray(
            shape=(num,),
            buffer=np.array([Reactor(kind, n_assems=n_assems) \
                                 for i in range(num)]),
            dtype=Reactor)
        return {k: gen_ary(k, self._rlztn.n_rxtrs[k], n_assems[k]) \
                    for k in rkinds}

    def _get_requesters(self):
        # requires self._rlztn to be set
        gen_ary = lambda kind, num: \
            np.ndarray(
            shape=(num,),
            buffer=np.array([Requester(kind, self.gids, self.nids) \
                                 for i in range(num)]),
            dtype=Requester)
        return {k: gen_ary(k, v) for k, v in self._rlztn.n_reqrs.items()}

    def _gen_structure(self, point, reactors, requesters):
        # requires self._rlztn to be set
        grps, nodes, arcs = [], [], []
        for rx_kind, rx_ary in reactors:
            for rxtr in rx_ary:
                for commod, nassems in self._rlztn.assem_dists[rx_kind].items():
                    for i in range(nassems):
                        excl_id = self.excl_ids.next()
                        gid = self.gids.next()
                        grp = rxtr.gen_group(gid)
                        grps.append(grp)
                        for reqr in requesters[commod]:
                            nid = self.nids.next()
                            node = rxtr.gen_node(nid, gid, excl_id)
                            arc = gen_arc(self.arcids.next(), point, commod,
                                          nid, rxtr, reqr)
                            nodes.append(node)
                            arcs.append(arc)
        return grps, nodes, arcs

    def gen_inst(self, point):
        """Parameters
        ----------
        point :  structured_species.Point
            A representation of a point in parameter space

        Returns
        -------
        inst : tuple of lists of ExGroups, ExNodes, and ExArgs
            A representation of a problem instance to be used by this species'
            family
        """
        # reset id generation
        self.nids = cyctools.Incrementer()
        self.excl_ids = cyctools.Incrementer()
        self.gids = cyctools.Incrementer()
        self.arcids = cyctools.Incrementer()

        # species objects
        if self._rlztn is None:
            # this could have been set before calling gen_inst, e.g., for
            # testing
            self._rlztn = pnt_to_realization(point)
        reactors = self._get_reactors()
        requesters = self._get_requesters()

        # structure
        rx_groups, rx_nodes, arcs = self._gen_structure(point, reactors, requesters)

        # combine groups, nodes
        groups = np.concatenate(
            (rx_groups,
             [x.group for ary in requesters.values() for x in ary]))
        nodes = np.concatenate(
            (rx_nodes,
             [x.nodes for ary in requesters.values() for x in ary]))

        return groups, nodes, arcs
<commit_msg>fixed some basic bugs, names of static functions, typos, etc<commit_after>"""A module for defining structured supply species.

:author: Matthew Gidden <matthew.gidden _at_ gmail.com>
"""
import itertools
import numpy as np
import random
import math

from collections import OrderedDict, defaultdict, Iterable, namedtuple

from cyclopts import tools as cyctools
from cyclopts import cyclopts_io as cycio
import cyclopts.exchange_instance as exinst
from cyclopts.problems import ProblemSpecies
from cyclopts.exchange_family import ResourceExchange

from cyclopts.structured_species import data
from cyclopts.structured_species import tools as strtools
from cyclopts.structured_species import request

class Point(strtools.Point):
    """A container class representing a point in parameter space"""

    """ordered mapping from input parameters to default values and np.dtypes, see
    the theory manual for further explanation of the parameter names"""
    parameters = OrderedDict(sorted(request.Point.parameters.items() + {
                "d_th": strtools.Param([0.67, 0.33, 0], (np.float64, 3)),
                "d_f_mox": strtools.Param([0., 0., 1., 0.], (np.float64, 4)),
                "d_f_thox": strtools.Param([0., 0., 0., 1.], (np.float64, 4)),
                "f_repo": strtools.Param(0.1, np.float32),
                }.items(), key=lambda t: t[0]))

    def __init__(self, d=None):
        """Parameters
        ----------
        d : dict, optional
            a dictionary with key value pairs of parameter name, parameter
            value
        """
        super(Point, self).__init__(d)
        if self.seed > 0:
            random.seed(self.seed)

    def _parameters(self):
        """ordered mapping from input parameters to default values and np.dtypes, see
        the theory manual for further explanation of the parameter names"""
        return Point.parameters

class Reactor(strtools.Reactor):
    """An extension reactor model for Structured Supply Species"""
    def __init__(self, kind, point=None, n_assems=None):
        super(Reactor, self).__init__(kind, point, n_assems)
        self.assem_qty = data.fuel_unit * data.core_vol_frac[self.kind] \
            / self.n_assems

    def gen_group(self, gid):
        supply = False
        grp = exinst.ExGroup(gid, supply)
        grp.AddCap(self.assem_qty)
        return grp

    def gen_node(self, nid, gid, excl_id):
        supply = False
        excl = True
        return exinst.ExNode(nid, gid, supply,
                             self.assem_qty, excl, excl_id)

class Requester(object):
    """A simplified requester model for Structured Supply Species"""
    def __init__(self, kind, gids, nids):
        self.kind = kind
        self.req_qty = data.sup_rhs[self.kind]
        gid = gids.next()
        req = True
        if self.kind == data.Supports.repo:
            self.group = grp = exinst.ExGroup(gid, req, self.req_qty)
        else:
            self.group = grp = exinst.ExGroup(gid, req, self.req_qty)
            commod = data.sup_to_commod[self.kind]
            rxtr = data.sup_to_rxtr[self.kind]
            grp.AddCap(self.req_qty * strtools.mean_enr(rxtr, commod) / 100 \
                           * data.relative_qtys[rxtr][commod])
        self._gen_nodes(gid, nids)
        self.loc = data.loc()

    def _gen_nodes(self, gid, nids):
        self.nodes = []
        self.commod_to_nodes = {}
        req = True
        for commod in data.sup_pref_basis[self.kind].keys():
            nid = nids.next()
            node = exinst.ExNode(nid, gid, req, self.req_qty)
            self.nodes.append(node)
            self.commod_to_nodes[commod] = node

    def coeff(self, enr, rkind, commod):
        if self.kind == data.Supports.repo:
            raise RuntimeError('Coeff not supported for repos')
        return enr / 100 * data.relative_qtys[rkind][commod]

class StructuredSupply(ProblemSpecies):
    """A class representing structured supply-based exchanges species."""

    @staticmethod
    def pnt_to_realization(point):
        """Returns a realization of a structured supply instance given a point
        in parameter space.

        A realization is a namedtuple of :
          * reqrs: a dictionary of the kind and number of each requester
          * rxtrs: a dictionary of the kind and number of each reactor
          * assem_dists: a dictionary of the kind of reactor to a dictionary
            of Commodity type to the number of assemblies of that Commodity type
        """
        reqrs = {data.Supports[i]: n \
                     for i, n in enumerate(strtools.support_breakdown(point))}
        rxtrs = {data.Reactors[i]: n \
                     for i, n in enumerate(strtools.reactor_breakdown(point))}
        dists = {k: strtools.assembly_breakdown(point, k) \
                     for k in data.Reactors}
        keys = ['n_reqrs', 'n_rxtrs', 'assem_dists']
        return namedtuple('Realization', keys)(reqrs, rxtrs, dists)

    @staticmethod
    def gen_arc(aid, point, commod, rx_node_id, rxtr, reqr):
        """generate an arc"""
        pref = strtools.preference(data.sup_pref_basis[reqr.kind][commod],
                                   rxtr.loc, reqr.loc,
                                   point.f_loc, point.r_l_c, point.n_reg)
        rq_coeffs = [reqr.coeff(rxtr.enr(commod), rxtr.kind, commod)] \
            if not reqr.kind == data.Supports.repo else []
        arc = exinst.ExArc(aid,
                           reqr.commod_to_nodes[commod].id, rq_coeffs,
                           rx_node_id, [1],
                           pref)
        return arc

    def __init__(self):
        super(StructuredSupply, self).__init__()
        self._family = ResourceExchange()
        self.space = None
        self._n_points = None
        # 16 bytes for uuid
        self.param_tbl_name = 'Points'
        self._param_dtype = np.dtype(
            [('paramid', ('str', 16)), ('family', ('str', 30))] + \
                [(name, param.dtype) for name, param in Point.parameters.items()])
        self.sum_tbl_name = 'Summary'
        facs = ['n_r_th', 'n_r_f_mox', 'n_r_f_thox', 'n_s_uox', 'n_s_th_mox',
                'n_s_f_mox', 'n_s_f_thox', 'n_s_repo']
        self.iter_params = ['d_th', 'd_f_mox', 'd_f_thox']
        self._sum_dtype = np.dtype(
            [('paramid', ('str', 16)), ('family', ('str', 30))] + \
                [(name, np.uint32) for name in facs])
        # reset id generation
        self.nids = cyctools.Incrementer()
        self.excl_ids = cyctools.Incrementer()
        self.gids = cyctools.Incrementer()
        self.arcids = cyctools.Incrementer()

        # default realization is None
        self._rlztn = None

    @property
    def family(self):
        """Returns
        -------
        family : ResourceExchange
            An instance of this species' family
        """
        return self._family

    @property
    def name(self):
        """Returns
        -------
        name : string
            The name of this species
        """
        return 'StructuredSupply'

    def register_tables(self, h5file, prefix):
        """Parameters
        ----------
        h5file : PyTables File
            the hdf5 file
        prefix : string
            the absolute path to the group for tables of this species

        Returns
        -------
        tables : list of cyclopts_io.Tables
            All tables that could be written to by this species.
        """
        return [cycio.Table(h5file, '/'.join([prefix, self.param_tbl_name]),
                            self._param_dtype),
                cycio.Table(h5file, '/'.join([prefix, self.sum_tbl_name]),
                            self._sum_dtype)]

    def read_space(self, space_dict):
        """Parameters
        ----------
        space_dict : dict
            A dictionary container resulting from the reading in of a run
            control file
        """
        self.space = {k: v if isinstance(v, Iterable) else [v] \
                          for k, v in space_dict.items() \
                          if k in Point.parameters}

    @property
    def n_points(self):
        """Returns
        -------
        n : int
            The total number of points in the parameter space
        """
        return cyctools.n_permutations(self.space, iter_keys=self.iter_params)

    def points(self):
        """Derived classes must implement this function returning a
        representation of a point in its parameter space to be used by other
        class member functions.

        Returns
        -------
        point_generator : generator
            A generator for representation of a point in parameter space to be
            used by this species
        """
        keys = self.space.keys()
        for k in keys:
            if k in self.iter_params:
                # iterable params must be iterable
                if not cyctools.seq_not_str(self.space[k]):
                    raise RuntimeError('{0} entry must be a Sequence'.format(k))
                # if they are defined as a single value, make them a sequence
                if not cyctools.seq_not_str(self.space[k][0]):
                    self.space[k] = [self.space[k]]
        vals = self.space.values()
        for args in cyctools.expand_args(vals, iter_keys=self.iter_params):
            d = {keys[i]: args[i] for i in range(len(args))}
            print(d)
            yield Point(d)

    def record_point(self, point, param_uuid, tables):
        """Parameters
        ----------
        point : tuple or other
            A representation of a point in parameter space
        param_uuid : uuid
            The uuid of the point in parameter space
        tables : list of cyclopts_io.Table
            The tables that can be written to
        """
        uid = param_uuid.bytes if len(param_uuid.bytes) == 16 \
            else param_uuid.bytes + '\0'

        data = [param_uuid.bytes, self._family.name]
        data += [getattr(point, k) for k in Point.parameters.keys()]
        tables[self.param_tbl_name].append_data([tuple(data)])

        data = [param_uuid.bytes, self._family.name]
        data += strtools.reactor_breakdown(point)
        data += strtools.support_breakdown(point)
        tables[self.sum_tbl_name].append_data([tuple(data)])

    def _get_reactors(self):
        # requires self._rlztn to be set
        rkinds = self._rlztn.n_rxtrs.keys()
        n_assems = {k: sum(v.values()) \
                       for k, v in self._rlztn.assem_dists.items()}
        gen_ary = lambda kind, num, n_assems: \
            np.ndarray(
            shape=(num,),
            buffer=np.array([Reactor(kind, n_assems=n_assems) \
                                 for i in range(num)]),
            dtype=Reactor)
        return {k: gen_ary(k, self._rlztn.n_rxtrs[k], n_assems[k]) \
                    for k in rkinds}

    def _get_requesters(self):
        # requires self._rlztn to be set
        gen_ary = lambda kind, num: \
            np.ndarray(
            shape=(num,),
            buffer=np.array([Requester(kind, self.gids, self.nids) \
                                 for i in range(num)]),
            dtype=Requester)
        return {k: gen_ary(k, v) for k, v in self._rlztn.n_reqrs.items()}

    def _gen_structure(self, point, reactors, requesters):
        # requires self._rlztn to be set
        grps, nodes, arcs = [], [], []
        for rx_kind, rx_ary in reactors:
            for rxtr in rx_ary:
                for commod, nassems in self._rlztn.assem_dists[rx_kind].items():
                    for i in range(nassems):
                        excl_id = self.excl_ids.next()
                        gid = self.gids.next()
                        grp = rxtr.gen_group(gid)
                        grps.append(grp)
                        for reqr in requesters[commod]:
                            nid = self.nids.next()
                            node = rxtr.gen_node(nid, gid, excl_id)
                            arc = gen_arc(self.arcids.next(), point, commod,
                                          nid, rxtr, reqr)
                            nodes.append(node)
                            arcs.append(arc)
        return grps, nodes, arcs

    def gen_inst(self, point):
        """Parameters
        ----------
        point :  structured_species.Point
            A representation of a point in parameter space

        Returns
        -------
        inst : tuple of lists of ExGroups, ExNodes, and ExArgs
            A representation of a problem instance to be used by this species'
            family
        """
        # reset id generation
        self.nids = cyctools.Incrementer()
        self.excl_ids = cyctools.Incrementer()
        self.gids = cyctools.Incrementer()
        self.arcids = cyctools.Incrementer()

        # species objects
        if self._rlztn is None:
            # this could have been set before calling gen_inst, e.g., for
            # testing
            self._rlztn = StructuredSupply.pnt_to_realization(point)
        reactors = self._get_reactors()
        requesters = self._get_requesters()

        # structure
        rx_groups, rx_nodes, arcs = self._gen_structure(point, reactors, requesters)

        # combine groups, nodes
        groups = np.concatenate(
            (rx_groups,
             [x.group for ary in requesters.values() for x in ary]))
        nodes = np.concatenate(
            (rx_nodes,
             [x.nodes for ary in requesters.values() for x in ary]))

        return groups, nodes, arcs
<|endoftext|>


<commit_before>import json

from future.utils import implements_iterator

from .document import Document, Edge
from .theExceptions import QueryError, AQLQueryError, SimpleQueryError, CreationError
from . import consts as CONST

__all__ = ["Query", "AQLQuery", "SimpleQuery", "Cursor", "RawCursor"]

@implements_iterator
class RawCursor(object) :
    "a raw interface to cursors that returns json"
    def __init__(self, database, cursorId) :
        self.database = database
        self.connection = self.database.connection
        self.id = cursorId
        self.URL = "%s/cursor/%s" % (self.database.URL, self.id)

    def __next__(self) :
        "returns the next batch"
        r = self.connection.session.put(self.URL)
        data = r.json()
        if r.status_code == 400 :
            raise CursorError(data["errorMessage"], self.id, data)
        return r.json()

@implements_iterator
class Query(object) :
    "This class is abstract and should not be instanciated. All query classes derive from it"

    def __init__(self, request, database, rawResults) :
        "If rawResults = True, the results will be returned as dictionaries instead of Document objects."

        self.rawResults = rawResults
        self.response = request.json()
        if self.response.get("error") and self.response["errorMessage"] != "no match" :
            raise QueryError(self.response["errorMessage"], self.response)

        self.request = request
        self.database = database
        self.connection = self.database.connection
        self.currI = 0
        if request.status_code == 201 or request.status_code == 200 or request.status_code == 202:
            self.batchNumber = 1
            try : #if there's only one element
                self.response = {"result" : [self.response["document"]], 'hasMore' : False}
                del(self.response["document"])
            except KeyError :
                pass

            if "hasMore" in self.response and self.response["hasMore"] :
                self.cursor = RawCursor(self.database, self.id)
            else :
                self.cursor = None
        elif request.status_code == 404 :
            self.batchNumber = 0
            self.result = []
        else :
            self._raiseInitFailed(request)

    def _raiseInitFailed(self, request) :
        "must be implemented in child, this called if the __init__ fails"
        raise NotImplemented("Must be implemented in child")

    def _developDoc(self, i) :
        """private function that transforms a json returned by ArangoDB into a pyArango Document or Edge"""
        docJson = self.result[i]
        try :
            collection = self.database[docJson["_id"].split("/")[0]]
        except KeyError :
            raise CreationError("result %d is not a valid Document. Try setting rawResults to True" % i)

        if collection.type == CONST.COLLECTION_EDGE_TYPE :
            self.result[i] = Edge(collection, docJson)
        else :
            self.result[i] = Document(collection, docJson)

    def nextBatch(self) :
        "become the next batch. raises a StopIteration if there is None"
        self.batchNumber += 1
        self.currI = 0
        try :
            if not self.response["hasMore"] or self.cursor is None :
                raise StopIteration("That was the last batch")
        except KeyError :
            raise AQLQueryError(self.response["errorMessage"], self.query, self.response)

        self.response = next(self.cursor)

    def delete(self) :
        "kills the cursor"
        self.connection.session.delete(self.cursor)

    def __next__(self) :
        """returns the next element of the query result. Automatomatically calls for new batches if needed"""
        try :
            v = self[self.currI]
        except IndexError :
            self.nextBatch()
        v = self[self.currI]
        self.currI += 1
        return v

    def __iter__(self) :
        """Returns an itererator so you can do::

            for doc in query : print doc
        """
        return self

    def __getitem__(self, i) :
        "returns a ith result of the query."
        if not self.rawResults and (self.result[i].__class__ is not Edge and self.result[i].__class__ is not Document) :
            self._developDoc(i)
        return self.result[i]

    def __len__(self) :
        """Returns the number of elements in the query results"""
        return len(self.result)

    def __getattr__(self, k) :
        try :
            resp = object.__getattribute__(self, "response")
            return resp[k]
        except (KeyError, AttributeError) :
            raise  AttributeError("There's no attribute %s" %(k))

    def __str__(self) :
        return str(self.result)

class AQLQuery(Query) :
    "AQL queries are attached to and instanciated by a database"
    def __init__(self, database, query, batchSize, bindVars, options, count, fullCount, rawResults = True,
                 json_encoder = None, **moreArgs) :
        # fullCount is passed in the options dict per https://docs.arangodb.com/3.1/HTTP/AqlQueryCursor/AccessingCursors.html
        options["fullCount"] = fullCount
        payload = {'query' : query, 'batchSize' : batchSize, 'bindVars' : bindVars, 'options' : options, 'count' : count}
        payload.update(moreArgs)

        self.query = query
        self.database = database
        self.connection = self.database.connection
        request = self.connection.session.post(database.cursorsURL, data = json.dumps(payload, cls=json_encoder))
        Query.__init__(self, request, database, rawResults)

    def explain(self, allPlans = False) :
        """Returns an explanation of the query. Setting allPlans to True will result in ArangoDB returning all possible plans. False returns only the optimal plan"""
        return self.database.explainAQLQuery(self.query, allPlans)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise AQLQueryError(data["errorMessage"], self.query, data)

class Cursor(Query) :
    "Cursor queries are attached to and instanciated by a database, use them to continue from where you left"
    def __init__(self, database, cursorId, rawResults) :
        self.rawResults = rawResults
        self._developed = set()
        self.batchNumber = 1
        self.cursor = RawCursor(database, cursorId)
        self.response = next(self.cursor)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise CursorError(data["errorMessage"], self.id, data)

class SimpleQuery(Query) :
    "Simple queries are attached to and instanciated by a collection"
    def __init__(self, collection, queryType, rawResults, json_encoder = None,
                 **queryArgs) :

        self.collection = collection
        self.connection = self.collection.database.connection

        payload = {'collection' : collection.name}
        payload.update(queryArgs)
        payload = json.dumps(payload, cls=json_encoder)
        URL = "%s/simple/%s" % (collection.database.URL, queryType)
        request = self.connection.session.put(URL, data = payload)

        Query.__init__(self, request, collection.database, rawResults)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise SimpleQueryError(data["errorMessage"], data)

    def _developDoc(self, i) :
        docJson = self.result[i]
        if self.collection.type == CONST.COLLECTION_EDGE_TYPE :
            self.result[i] = Edge(self.collection, docJson)
        else :
            self.result[i] = Document(self.collection, docJson)
<commit_msg>Correctly read and populate the cursor id for queries<commit_after>import json

from future.utils import implements_iterator

from .document import Document, Edge
from .theExceptions import QueryError, AQLQueryError, SimpleQueryError, CreationError
from . import consts as CONST

__all__ = ["Query", "AQLQuery", "SimpleQuery", "Cursor", "RawCursor"]

@implements_iterator
class RawCursor(object) :
    "a raw interface to cursors that returns json"
    def __init__(self, database, cursorId) :
        self.database = database
        self.connection = self.database.connection
        self.id = cursorId
        self.URL = "%s/cursor/%s" % (self.database.URL, self.id)

    def __next__(self) :
        "returns the next batch"
        r = self.connection.session.put(self.URL)
        data = r.json()
        if r.status_code == 400 :
            raise CursorError(data["errorMessage"], self.id, data)
        return r.json()

@implements_iterator
class Query(object) :
    "This class is abstract and should not be instanciated. All query classes derive from it"

    def __init__(self, request, database, rawResults) :
        "If rawResults = True, the results will be returned as dictionaries instead of Document objects."

        self.rawResults = rawResults
        self.response = request.json()
        if self.response.get("error") and self.response["errorMessage"] != "no match" :
            raise QueryError(self.response["errorMessage"], self.response)

        self.request = request
        self.database = database
        self.connection = self.database.connection
        self.currI = 0
        if request.status_code == 201 or request.status_code == 200 or request.status_code == 202:
            self.batchNumber = 1
            try : #if there's only one element
                self.response = {"result" : [self.response["document"]], 'hasMore' : False}
                del(self.response["document"])
            except KeyError :
                pass

            if "hasMore" in self.response and self.response["hasMore"] :
		cursor_id = self.response.get("id","")
                self.cursor = RawCursor(self.database, cursor_id)
            else :
                self.cursor = None
        elif request.status_code == 404 :
            self.batchNumber = 0
            self.result = []
        else :
            self._raiseInitFailed(request)

    def _raiseInitFailed(self, request) :
        "must be implemented in child, this called if the __init__ fails"
        raise NotImplemented("Must be implemented in child")

    def _developDoc(self, i) :
        """private function that transforms a json returned by ArangoDB into a pyArango Document or Edge"""
        docJson = self.result[i]
        try :
            collection = self.database[docJson["_id"].split("/")[0]]
        except KeyError :
            raise CreationError("result %d is not a valid Document. Try setting rawResults to True" % i)

        if collection.type == CONST.COLLECTION_EDGE_TYPE :
            self.result[i] = Edge(collection, docJson)
        else :
            self.result[i] = Document(collection, docJson)

    def nextBatch(self) :
        "become the next batch. raises a StopIteration if there is None"
        self.batchNumber += 1
        self.currI = 0
        try :
            if not self.response["hasMore"] or self.cursor is None :
                raise StopIteration("That was the last batch")
        except KeyError :
            raise AQLQueryError(self.response["errorMessage"], self.query, self.response)

        self.response = next(self.cursor)

    def delete(self) :
        "kills the cursor"
        self.connection.session.delete(self.cursor)

    def __next__(self) :
        """returns the next element of the query result. Automatomatically calls for new batches if needed"""
        try :
            v = self[self.currI]
        except IndexError :
            self.nextBatch()
        v = self[self.currI]
        self.currI += 1
        return v

    def __iter__(self) :
        """Returns an itererator so you can do::

            for doc in query : print doc
        """
        return self

    def __getitem__(self, i) :
        "returns a ith result of the query."
        if not self.rawResults and (self.result[i].__class__ is not Edge and self.result[i].__class__ is not Document) :
            self._developDoc(i)
        return self.result[i]

    def __len__(self) :
        """Returns the number of elements in the query results"""
        return len(self.result)

    def __getattr__(self, k) :
        try :
            resp = object.__getattribute__(self, "response")
            return resp[k]
        except (KeyError, AttributeError) :
            raise  AttributeError("There's no attribute %s" %(k))

    def __str__(self) :
        return str(self.result)

class AQLQuery(Query) :
    "AQL queries are attached to and instanciated by a database"
    def __init__(self, database, query, batchSize, bindVars, options, count, fullCount, rawResults = True,
                 json_encoder = None, **moreArgs) :
        # fullCount is passed in the options dict per https://docs.arangodb.com/3.1/HTTP/AqlQueryCursor/AccessingCursors.html
        options["fullCount"] = fullCount
        payload = {'query' : query, 'batchSize' : batchSize, 'bindVars' : bindVars, 'options' : options, 'count' : count}
        payload.update(moreArgs)

        self.query = query
        self.database = database
        self.connection = self.database.connection
        request = self.connection.session.post(database.cursorsURL, data = json.dumps(payload, cls=json_encoder))
        Query.__init__(self, request, database, rawResults)

    def explain(self, allPlans = False) :
        """Returns an explanation of the query. Setting allPlans to True will result in ArangoDB returning all possible plans. False returns only the optimal plan"""
        return self.database.explainAQLQuery(self.query, allPlans)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise AQLQueryError(data["errorMessage"], self.query, data)

class Cursor(Query) :
    "Cursor queries are attached to and instanciated by a database, use them to continue from where you left"
    def __init__(self, database, cursorId, rawResults) :
        self.rawResults = rawResults
        self._developed = set()
        self.batchNumber = 1
        self.cursor = RawCursor(database, cursorId)
        self.response = next(self.cursor)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise CursorError(data["errorMessage"], self.id, data)

class SimpleQuery(Query) :
    "Simple queries are attached to and instanciated by a collection"
    def __init__(self, collection, queryType, rawResults, json_encoder = None,
                 **queryArgs) :

        self.collection = collection
        self.connection = self.collection.database.connection

        payload = {'collection' : collection.name}
        payload.update(queryArgs)
        payload = json.dumps(payload, cls=json_encoder)
        URL = "%s/simple/%s" % (collection.database.URL, queryType)
        request = self.connection.session.put(URL, data = payload)

        Query.__init__(self, request, collection.database, rawResults)

    def _raiseInitFailed(self, request) :
        data = request.json()
        raise SimpleQueryError(data["errorMessage"], data)

    def _developDoc(self, i) :
        docJson = self.result[i]
        if self.collection.type == CONST.COLLECTION_EDGE_TYPE :
            self.result[i] = Edge(self.collection, docJson)
        else :
            self.result[i] = Document(self.collection, docJson)
<|endoftext|>