<commit_before>""" Tools for creating and manipulating animations

Animations are objects created from classes derived from LSAnimation. Each
LSAnimation object maintains a list of frames which can be manipulated using
methods specific to the subclass. Each frame is a list of integers representing
segments that need to be illuminated. The first item in the list represents the
width of the frame in columns. All of the frames that make up a given animation
must have the same dimensions.

The rest of a the numbers in the frame object are a repeating pattern of Red,
Green, and Blue color masks, every three belonging to tiles succesively scanned
from the left to the right and the top to the bottom. A color mask is an integer
defining a certain shape on the 7 segment display. The shapes referenced in Red,
Green, and Blue will be composited over eachother on the display. Look at the
documentation in the Shapes module for more information on the scheme used to
encode different shapes.

Frames can be constructed and manipulated using the LSFrameGen class. Objects
created by LSFrameGen maintain an internal representation of a frame and output
the compacted list-type frames required by LSAnimation whenever the get() method
is called. A simple technique for programatically generating animations is thus
to maintain one or more LSFrameGen objects that you apply transformations to
while iteratively adding the output of get() to your LSAnimation object.

Using the mergeFrames tool you can merge the contents of a frame, at an
arbitrary offset, into another frame of equal or greater size. By unpacking and
merging the frames from two animations you can create a new combined animation.

Example:

    # Import the LightSweeper API
    from lightsweeper.lsapi import *

    # Also import the animation library
    from lightsweeper import lsanimate

    # Create a LightSweeper display:
    ourDisplay = LSDisplay()

    # We'll start by creating a frame the size of our desired animation
    # Let's make this one fill the entire board
    frame = lsanimate.LSFrameGen(ourDisplay.rows, ourDisplay.cols)

    # We'll also need a fresh animation object to put our generated frames in
    ourAnimation = lsanimate.LSAnimation()

    # We want the background of this animation to be all pink so we need to
    # build a colormask that looks like a pink eight. Pink is red+blue so:
    pinkEight = (redMask, greenMask, blueMask) = (Shapes.EIGHT, 0, Shapes.EIGHT)

    # Then fill the frame with it
    frame.fill(pinkEight)

    # For this animation we'd like to make a vertical green stripe that
    # moves down the the height of the frame. We'll define three new masks:
    bottomLine = (Shapes.A, Shapes.UNDERSCORE, Shapes.A)
    midLine = (Shapes.ZERO, Shapes.DASH, Shapes.ZERO)
    topLine = (Shapes.U+Shapes.DASH, Shapes.SEG_A, Shapes.U+Shapes.DASH)

    # Now we programatically animate the scanning line:
    for _ in range(10):      # Loop the animation 10 times
        for thisRow in range(frame.rows):
            for i in range(4):
                for thisCol in range(frame.cols):
                    # We just re-use the same LSFrameGen object, making incremental
                    # changes to it and adding the modified frames to our animation
                    if i is 0:
                        frame.edit(thisRow, thisCol, topLine)
                    elif i is 1:
                        frame.edit(thisRow, thisCol, midLine)
                    elif i is 2:
                        frame.edit(thisRow, thisCol, bottomLine)
                    else:
                        frame.edit(thisRow, thisCol, pinkEight)

                # Filter out frames that are all pink:
                if not all(i in [127, 0] for i in frame.get()[1:]):
                    ourAnimation.addFrame(frame.get())  # Add the constructed frame

    # Our animation is built! Now let's play it:
    ourAnimation.play(ourDisplay, frameRate=10)     # Set the frameRate to 0 to step
                                                    # through the animation manually

"""

from collections import defaultdict
import time
import itertools

from lightsweeper import Shapes
from lightsweeper import Colors
from lightsweeper import lsdisplay

# TODO: Custom error handlers

def validateFrame(frame):
   # print(repr(frame)) # Debugging
    frameLen = len(frame) - 1
    if frameLen % 3 is not 0:
        return False
    cells = frameLen/3
    if (cells/frame[0]).is_integer() is False:
        return False
    if all(i <= 128 for i in frame[1:]) is False:
        return False
    return True

def mergeFrames(baseFrame, subFrame, offset=(0,0)):
    # TODO: Alpha channel
    rows = bRows = frameRows(baseFrame)
    cols = bCols = baseFrame.pop(0)
    sRows = frameRows(subFrame)
    sCols = subFrame.pop(0)
    if sRows+offset[0] > rows or sCols+offset[1] > cols:
        raise Exception("Cannot merge frames, subFrame plus offset must be smaller than baseFrame")
    out = list()
    hiddenLayer = list()
    out.append(cols)
    out.extend([baseFrame.pop(0) for _ in range(offset[0]*bCols*3)])
    while len(subFrame) > 0:
        out.extend([baseFrame.pop(0) for _ in range(offset[1]*3)])
        hiddenLayer.extend([baseFrame.pop(0) for _ in range(sCols*3)])
        out.extend([subFrame.pop(0) for _ in range(sCols*3)])
        out.extend([baseFrame.pop(0) for _ in range((bCols-sCols-offset[1])*3)])
    out.extend(baseFrame)
    return(out)

def frameRows (frame):
    return(len(frame[1:])/frame[0]/3)

def frameCols (frame):
    return(frame[0])

class LSAnimation:

    def __init__(self):
        self._frames = list()

    def numFrames(self):
        return(len(self._frames))

    def addFrame(self, frame):
        if len(self._frames) is 0:
            self.rows = int(len(frame[1:])/frame[0]/3)
            self.cols = frame[0]
        self.insertFrame(len(self._frames), frame)
        return True

    def insertFrame(self, index, frame):
    # Inserts before given frame
        if validateFrame(frame) is False:
            print("Error frame is invalid")
            raise Exception
        if self._checkIndex(index) is True:
            self._frames.insert(index,frame)
        return True

    def deleteFrame(self, index):
        index -= 1
        if self._checkIndex(index) is True:
            frames = self._frames[0:index]
            frames.extend(self._frames[index+1:])
            self._frames = frames
        return True

    def dropFrames(self):
        # Discard all stored frames
        self._frames = list()

    def showFrames(self):
        Frame = self.nextFrame()
        i = 0
        while True:
            i += 1
            try:
                print("Frame {:d}: {:s}".format(i, str(next(Frame))))
            except StopIteration:
                print("Out of frames")
                break

    def nextFrame(self):
        if len(self._frames) is 0:
            print("No frames!")
        else:
            for frame in self._frames:
                yield(frame)

    def play(self, display, frameRate = 30):
        if frameRows(self._frames[0]) > display.rows or frameCols(self._frames[0]) > display.cols:
            raise Exception("Animation is too large for this display!")  # TODO: support clipping
        if frameRate < 0:
            raise ValueError("Please specify a non-negative frame rate")
        print("Starting animation ({:d} frames at {:d} fps)".format(len(self._frames), frameRate))
        Frame = self.nextFrame()
        i = 0
        for frame in Frame:
            i+=1
            if frameRate is 0:
                input("Press any key to advance the animation...\n")
                print("Frame {:d}: {:s}".format(i, str(frame)))
            stime = time.time()
            display.floor.renderFrame(frame)
            display.heartbeat()
            renderTime = time.time() - stime

            if frameRate is not 0:
                fps = 1.0/renderTime
                if fps < frameRate:
                    print("[Animation]{0:.4f} FPS".format(1.0/renderTime), end="\r")
                else:
                    time.sleep((1.0/frameRate) - renderTime)
        if frameRate is 0:
            print("Animation ended.")
        print(" " * 22, end = "\r")

    def _checkIndex(self, index):
        if index not in range(-1,len(self._frames)+1):
            print("Error frame index out of range")
            raise Exception
        return True

class ScrollingText(LSAnimation):

    height = 1
    width = 1
    color = Colors.WHITE
    direction = "left"
    iterations = 1

    def __init__(self, text, color=Colors.WHITE, height=1, width=None):
        super().__init__()
        self.charString = Shapes.stringToShapes(text)
        self.direction="left"
        self.height = height
        self.color = color
        if width is None:
            self.width = len(self.charString)
        else:
            self.width = width

    def __setattr__ (self, name, value):
        super().__setattr__(name, value)
        if name in ["height", "width", "color", "direction", "iterations"]:
            self._buildAnimation()

    def _buildAnimation (self):
        self.dropFrames()
        self.frame = LSFrameGen(self.height, self.width)
        cs = self.charString[:]
        if "right" in self.direction.lower():
            cs = reversed(cs)
        if self.height is 1:
            endCap = lambda x: self.width-1 if "left" in x.lower() else 0
            for i in range(0,self.iterations):
                for char in cs:
                    try:
                        color = next(self.color)
                    except TypeError:
                        color = self.color
                    colorMask = Colors.intToRGB(color)
                    charR = charG = charB = 0
                    (charR, charG, charB) = [char if i is not 0 else 0 for i in colorMask]
                    self.frame.edit(0, endCap(self.direction), (charR, charG, charB))
                    self.addFrame(self.frame.get())
                    self._pickShift()
                if True:
                    for i in range(0,self.width):
                        self.frame.edit(0, endCap(self.direction), (0, 0, 0))
                        self.addFrame(self.frame.get())
                        self._pickShift()
        else:
            raise NotImplementedError("height cannot be more than 1")

    def _pickShift(self):
        lowerD = self.direction.lower()
        if "left" in lowerD:
            self.frame.shiftLeft()
        elif "right" in lowerD:
            self.frame.shiftRight()

class LSFrameGen:

    def __init__(self, rows, cols):
        self.rows = rows
        self.cols = cols
        self.frame = defaultdict(lambda: defaultdict(int))

    # Allows you to edit an existing frame structure, if no colormask is set
    # then the tile will keep its current colormask
    def edit(self, row, col, colormask):
        if row > self.rows or col > self.cols:
            print("Edit error, index out of range")
            raise Exception
        self.frame[row][col] = [colormask]

    def fill(self, colormask):
        for row in range(0, self.rows):
            for col in range(0, self.cols):
                self.frame[row][col] = [colormask]

    def shiftLeft (self):
        for col in range(0, self.cols-1):
            for row in range(0, self.rows):
                self.frame[row][col] = self.frame[row][col+1]
        for row in range(0, self.rows):
            self.frame[row][self.cols-1] = [(0,0,0)]

    def shiftRight (self):
        for col in reversed(range(0, self.cols)):
            for row in range(0, self.rows):
                self.frame[row][col] = self.frame[row][col-1]
        for row in range(0, self.rows):
            self.frame[row][0] = [(0,0,0)]

    def shiftUp (self):
        for row in range(0, self.rows-1):
            for col in range(0, self.cols):
                self.frame[row][col] = self.frame[row+1][col]
        for col in range(0, self.cols):
            self.frame[self.rows-1][col] = [(0,0,0)]

    def shiftDown (self):
        for row in reversed(range(0, self.rows)):
            for col in range(0, self.cols):
                self.frame[row][col] = self.frame[row-1][col]
        for col in range(0, self.cols):
            self.frame[0][col] = [(0,0,0)]

#TODO: (shiftDiag?)

    def print(self):
        for row in self.frame:
            print([self.frame[row][col] for col in self.frame[row]])

    def get(self):
        frameOut = list()
        frameOut.append(self.cols)
        for row in range(0, self.rows):
            for col in range(0, self.cols):
                try:
                    cell = self.frame[row][col][0]
                    frameOut.append(cell[0])
                    frameOut.append(cell[1])
                    frameOut.append(cell[2])
                except:
                #    print("Warning: ({:d},{:d}) has no update".format(row,col)) # Debugging
                    frameOut.append(128)
                    frameOut.append(128)
                    frameOut.append(128)
        return(frameOut)  # frameOut is a list consisting of the number of columns in the frame
                          # followed by a repeating pattern of 3 integers, each representing a
                          # subsequent tile's red, green, and blue colormasks

  #  def put(self, frame):
  #      if not validateFrame(frame):
  #          raise Exception("Cannot put frame. Frame is invalid")
  #      cols = frame.pop()
  #      row = col = 0
  #      for _ in itertools.repeat(None, int(len(frame)/3)):
  #          if col is cols:
  #              row += 1
  #              col = 0
  #          rMask, gMask, bMask = frame.pop(0), frame.pop(0), frame.pop(0)
  #          if rMask is not 128:
  #              self.frame[row][col] = (rMask, gMask, bMask)
  #          col += 1

def main():
    print("Importing LSDisplay")
    import lsdisplay

    d = lsdisplay.LSDisplay(initScreen=False)

    colormask = (Shapes.ZERO, 0,0)
    diffmask = (Shapes.ONE, Shapes.TWO, Shapes.THREE)
    redZero = (Shapes.ZERO, Shapes.OFF, Shapes.OFF)
    greenZero = (0, Shapes.ZERO, 0)
    blueZero = (0, 0, 126)

    ourAnimation = LSAnimation()

    frame = LSFrameGen(3,d.cols)

    for _ in range(0,100):
        for i in range(d.cols):
            frame.edit(0,i,redZero)
            frame.edit(1,i,greenZero)
            frame.edit(2,i,blueZero)
        ourAnimation.addFrame(frame.get())
        for i in range(0,d.cols):
            frame.edit(1,i,redZero)
            frame.edit(2,i,greenZero)
            frame.edit(0,i,blueZero)
        ourAnimation.addFrame(frame.get())
        for i in range(0,d.cols):
            frame.edit(2,i,redZero)
            frame.edit(0,i,greenZero)
            frame.edit(1,i,blueZero)
        ourAnimation.addFrame(frame.get())

   # ourAnimation.deleteFrame(1)

  #  ourAnimation.showFrames()

    ourAnimation.play(d, frameRate=15)

if __name__ == '__main__':
    main()
<commit_msg>Minor tweak to example<commit_after>""" Tools for creating and manipulating animations

Animations are objects created from classes derived from LSAnimation. Each
LSAnimation object maintains a list of frames which can be manipulated using
methods specific to the subclass. Each frame is a list of integers representing
segments that need to be illuminated. The first item in the list represents the
width of the frame in columns. All of the frames that make up a given animation
must have the same dimensions.

The rest of a the numbers in the frame object are a repeating pattern of Red,
Green, and Blue color masks, every three belonging to tiles succesively scanned
from the left to the right and the top to the bottom. A color mask is an integer
defining a certain shape on the 7 segment display. The shapes referenced in Red,
Green, and Blue will be composited over eachother on the display. Look at the
documentation in the Shapes module for more information on the scheme used to
encode different shapes.

Frames can be constructed and manipulated using the LSFrameGen class. Objects
created by LSFrameGen maintain an internal representation of a frame and output
the compacted list-type frames required by LSAnimation whenever the get() method
is called. A simple technique for programatically generating animations is thus
to maintain one or more LSFrameGen objects that you apply transformations to
while iteratively adding the output of get() to your LSAnimation object.

Using the mergeFrames tool you can merge the contents of a frame, at an
arbitrary offset, into another frame of equal or greater size. By unpacking and
merging the frames from two animations you can create a new combined animation.

Example:

    # Import the LightSweeper API
    from lightsweeper.lsapi import *

    # Also import the animation library
    from lightsweeper import lsanimate

    # Create a LightSweeper display:
    ourDisplay = LSDisplay()

    # We'll start by creating a frame the size of our desired animation
    # Let's make this one fill the entire board
    frame = lsanimate.LSFrameGen(ourDisplay.rows, ourDisplay.cols)

    # We'll also need a fresh animation object to put our generated frames in
    ourAnimation = lsanimate.LSAnimation()

    # We want the background of this animation to be all pink so we need to
    # build a colormask that looks like a pink eight. Pink is red+blue so:
    pinkEight = (redMask, greenMask, blueMask) = (Shapes.EIGHT, 0, Shapes.EIGHT)

    # Then fill the frame with it
    frame.fill(pinkEight)

    # For this animation we'd like to make a vertical green stripe that
    # moves down the the height of the frame. We'll define three new masks:
    bottomLine = (Shapes.A, Shapes.UNDERSCORE, Shapes.A)
    midLine = (Shapes.ZERO, Shapes.DASH, Shapes.ZERO)
    topLine = (Shapes.U+Shapes.DASH, Shapes.SEG_A, Shapes.U+Shapes.DASH)

    # Now we programatically animate the scanning line:
    for _ in range(10):      # Loop the animation 10 times
        for thisRow in range(frame.rows):
            for i in range(4):  # It takes 4 frames to move the line from the
                                # top of the figure 8 to the bottom and then
                                # reset the row
                for thisCol in range(frame.cols):
                    # We just re-use the same LSFrameGen object, making incremental
                    # changes to it and adding the modified frames to our animation
                    if i is 0:
                        frame.edit(thisRow, thisCol, topLine)
                    elif i is 1:
                        frame.edit(thisRow, thisCol, midLine)
                    elif i is 2:
                        frame.edit(thisRow, thisCol, bottomLine)
                    else:
                        frame.edit(thisRow, thisCol, pinkEight)

                # Filter out frames that are all pink:
                if not all(i in pinkEight for i in frame.get()[1:]):
                    ourAnimation.addFrame(frame.get())  # Add the constructed frame

    # Our animation is built! Now let's play it:
    ourAnimation.play(ourDisplay, frameRate=10)     # Set the frameRate to 0 to step
                                                    # through the animation manually

"""

from collections import defaultdict
import time
import itertools

from lightsweeper import Shapes
from lightsweeper import Colors
from lightsweeper import lsdisplay

# TODO: Custom error handlers

def validateFrame(frame):
   # print(repr(frame)) # Debugging
    frameLen = len(frame) - 1
    if frameLen % 3 is not 0:
        return False
    cells = frameLen/3
    if (cells/frame[0]).is_integer() is False:
        return False
    if all(i <= 128 for i in frame[1:]) is False:
        return False
    return True

def mergeFrames(baseFrame, subFrame, offset=(0,0)):
    # TODO: Alpha channel
    rows = bRows = frameRows(baseFrame)
    cols = bCols = baseFrame.pop(0)
    sRows = frameRows(subFrame)
    sCols = subFrame.pop(0)
    if sRows+offset[0] > rows or sCols+offset[1] > cols:
        raise Exception("Cannot merge frames, subFrame plus offset must be smaller than baseFrame")
    out = list()
    hiddenLayer = list()
    out.append(cols)
    out.extend([baseFrame.pop(0) for _ in range(offset[0]*bCols*3)])
    while len(subFrame) > 0:
        out.extend([baseFrame.pop(0) for _ in range(offset[1]*3)])
        hiddenLayer.extend([baseFrame.pop(0) for _ in range(sCols*3)])
        out.extend([subFrame.pop(0) for _ in range(sCols*3)])
        out.extend([baseFrame.pop(0) for _ in range((bCols-sCols-offset[1])*3)])
    out.extend(baseFrame)
    return(out)

def frameRows (frame):
    return(len(frame[1:])/frame[0]/3)

def frameCols (frame):
    return(frame[0])

class LSAnimation:

    def __init__(self):
        self._frames = list()

    def numFrames(self):
        return(len(self._frames))

    def addFrame(self, frame):
        if len(self._frames) is 0:
            self.rows = int(len(frame[1:])/frame[0]/3)
            self.cols = frame[0]
        self.insertFrame(len(self._frames), frame)
        return True

    def insertFrame(self, index, frame):
    # Inserts before given frame
        if validateFrame(frame) is False:
            print("Error frame is invalid")
            raise Exception
        if self._checkIndex(index) is True:
            self._frames.insert(index,frame)
        return True

    def deleteFrame(self, index):
        index -= 1
        if self._checkIndex(index) is True:
            frames = self._frames[0:index]
            frames.extend(self._frames[index+1:])
            self._frames = frames
        return True

    def dropFrames(self):
        # Discard all stored frames
        self._frames = list()

    def showFrames(self):
        Frame = self.nextFrame()
        i = 0
        while True:
            i += 1
            try:
                print("Frame {:d}: {:s}".format(i, str(next(Frame))))
            except StopIteration:
                print("Out of frames")
                break

    def nextFrame(self):
        if len(self._frames) is 0:
            print("No frames!")
        else:
            for frame in self._frames:
                yield(frame)

    def play(self, display, frameRate = 30):
        if frameRows(self._frames[0]) > display.rows or frameCols(self._frames[0]) > display.cols:
            raise Exception("Animation is too large for this display!")  # TODO: support clipping
        if frameRate < 0:
            raise ValueError("Please specify a non-negative frame rate")
        print("Starting animation ({:d} frames at {:d} fps)".format(len(self._frames), frameRate))
        Frame = self.nextFrame()
        i = 0
        for frame in Frame:
            i+=1
            if frameRate is 0:
                input("Press any key to advance the animation...\n")
                print("Frame {:d}: {:s}".format(i, str(frame)))
            stime = time.time()
            display.floor.renderFrame(frame)
            display.heartbeat()
            renderTime = time.time() - stime

            if frameRate is not 0:
                fps = 1.0/renderTime
                if fps < frameRate:
                    print("[Animation]{0:.4f} FPS".format(1.0/renderTime), end="\r")
                else:
                    time.sleep((1.0/frameRate) - renderTime)
        if frameRate is 0:
            print("Animation ended.")
        print(" " * 22, end = "\r")

    def _checkIndex(self, index):
        if index not in range(-1,len(self._frames)+1):
            print("Error frame index out of range")
            raise Exception
        return True

class ScrollingText(LSAnimation):

    height = 1
    width = 1
    color = Colors.WHITE
    direction = "left"
    iterations = 1

    def __init__(self, text, color=Colors.WHITE, height=1, width=None):
        super().__init__()
        self.charString = Shapes.stringToShapes(text)
        self.direction="left"
        self.height = height
        self.color = color
        if width is None:
            self.width = len(self.charString)
        else:
            self.width = width

    def __setattr__ (self, name, value):
        super().__setattr__(name, value)
        if name in ["height", "width", "color", "direction", "iterations"]:
            self._buildAnimation()

    def _buildAnimation (self):
        self.dropFrames()
        self.frame = LSFrameGen(self.height, self.width)
        cs = self.charString[:]
        if "right" in self.direction.lower():
            cs = reversed(cs)
        if self.height is 1:
            endCap = lambda x: self.width-1 if "left" in x.lower() else 0
            for i in range(0,self.iterations):
                for char in cs:
                    try:
                        color = next(self.color)
                    except TypeError:
                        color = self.color
                    colorMask = Colors.intToRGB(color)
                    charR = charG = charB = 0
                    (charR, charG, charB) = [char if i is not 0 else 0 for i in colorMask]
                    self.frame.edit(0, endCap(self.direction), (charR, charG, charB))
                    self.addFrame(self.frame.get())
                    self._pickShift()
                if True:
                    for i in range(0,self.width):
                        self.frame.edit(0, endCap(self.direction), (0, 0, 0))
                        self.addFrame(self.frame.get())
                        self._pickShift()
        else:
            raise NotImplementedError("height cannot be more than 1")

    def _pickShift(self):
        lowerD = self.direction.lower()
        if "left" in lowerD:
            self.frame.shiftLeft()
        elif "right" in lowerD:
            self.frame.shiftRight()

class LSFrameGen:

    def __init__(self, rows, cols):
        self.rows = rows
        self.cols = cols
        self.frame = defaultdict(lambda: defaultdict(int))

    # Allows you to edit an existing frame structure, if no colormask is set
    # then the tile will keep its current colormask
    def edit(self, row, col, colormask):
        if row > self.rows or col > self.cols:
            print("Edit error, index out of range")
            raise Exception
        self.frame[row][col] = [colormask]

    def fill(self, colormask):
        for row in range(0, self.rows):
            for col in range(0, self.cols):
                self.frame[row][col] = [colormask]

    def shiftLeft (self):
        for col in range(0, self.cols-1):
            for row in range(0, self.rows):
                self.frame[row][col] = self.frame[row][col+1]
        for row in range(0, self.rows):
            self.frame[row][self.cols-1] = [(0,0,0)]

    def shiftRight (self):
        for col in reversed(range(0, self.cols)):
            for row in range(0, self.rows):
                self.frame[row][col] = self.frame[row][col-1]
        for row in range(0, self.rows):
            self.frame[row][0] = [(0,0,0)]

    def shiftUp (self):
        for row in range(0, self.rows-1):
            for col in range(0, self.cols):
                self.frame[row][col] = self.frame[row+1][col]
        for col in range(0, self.cols):
            self.frame[self.rows-1][col] = [(0,0,0)]

    def shiftDown (self):
        for row in reversed(range(0, self.rows)):
            for col in range(0, self.cols):
                self.frame[row][col] = self.frame[row-1][col]
        for col in range(0, self.cols):
            self.frame[0][col] = [(0,0,0)]

#TODO: (shiftDiag?)

    def print(self):
        for row in self.frame:
            print([self.frame[row][col] for col in self.frame[row]])

    def get(self):
        frameOut = list()
        frameOut.append(self.cols)
        for row in range(0, self.rows):
            for col in range(0, self.cols):
                try:
                    cell = self.frame[row][col][0]
                    frameOut.append(cell[0])
                    frameOut.append(cell[1])
                    frameOut.append(cell[2])
                except:
                #    print("Warning: ({:d},{:d}) has no update".format(row,col)) # Debugging
                    frameOut.append(128)
                    frameOut.append(128)
                    frameOut.append(128)
        return(frameOut)  # frameOut is a list consisting of the number of columns in the frame
                          # followed by a repeating pattern of 3 integers, each representing a
                          # subsequent tile's red, green, and blue colormasks

  #  def put(self, frame):
  #      if not validateFrame(frame):
  #          raise Exception("Cannot put frame. Frame is invalid")
  #      cols = frame.pop()
  #      row = col = 0
  #      for _ in itertools.repeat(None, int(len(frame)/3)):
  #          if col is cols:
  #              row += 1
  #              col = 0
  #          rMask, gMask, bMask = frame.pop(0), frame.pop(0), frame.pop(0)
  #          if rMask is not 128:
  #              self.frame[row][col] = (rMask, gMask, bMask)
  #          col += 1

def main():
    print("Importing LSDisplay")
    import lsdisplay

    d = lsdisplay.LSDisplay(initScreen=False)

    colormask = (Shapes.ZERO, 0,0)
    diffmask = (Shapes.ONE, Shapes.TWO, Shapes.THREE)
    redZero = (Shapes.ZERO, Shapes.OFF, Shapes.OFF)
    greenZero = (0, Shapes.ZERO, 0)
    blueZero = (0, 0, 126)

    ourAnimation = LSAnimation()

    frame = LSFrameGen(3,d.cols)

    for _ in range(0,100):
        for i in range(d.cols):
            frame.edit(0,i,redZero)
            frame.edit(1,i,greenZero)
            frame.edit(2,i,blueZero)
        ourAnimation.addFrame(frame.get())
        for i in range(0,d.cols):
            frame.edit(1,i,redZero)
            frame.edit(2,i,greenZero)
            frame.edit(0,i,blueZero)
        ourAnimation.addFrame(frame.get())
        for i in range(0,d.cols):
            frame.edit(2,i,redZero)
            frame.edit(0,i,greenZero)
            frame.edit(1,i,blueZero)
        ourAnimation.addFrame(frame.get())

   # ourAnimation.deleteFrame(1)

  #  ourAnimation.showFrames()

    ourAnimation.play(d, frameRate=15)

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


<commit_before>'''
Reader for LSP output xdr files (.p4's)

'''
import xdrlib as xdr;
import itertools as itools;
import multiprocessing;
import numpy as np;

def make_points(d):
    d['x']=[]; d['y']=[]; d['z']=[];
    tmp1 = len(d['xp']);
    tmp2 = tmp1 * len(d['yp']);
    for i in range(d['nAll']):
        d['z'].append(d['zp'][i/tmp2]);
        d['y'].append(d['yp'][i%tmp2/tmp1]);
        d['x'].append(d['xp'][i%tmp1]);
    del d['zp'],d['yp'],d['xp'];
    return d;

def read_fields(d):
    for quantity in d['dqs']:
        d[quantity+'x']=[];
        d[quantity+'y']=[];
        d[quantity+'z']=[];
        for k in xrange(d['nAll']):
            d[quantity+'x'].append(d['x'+quantity].unpack_float());
            d[quantity+'y'].append(d['x'+quantity].unpack_float());
            d[quantity+'z'].append(d['x'+quantity].unpack_float());
        del d['x'+quantity];
    del d['dqs'],d['nAll'];
    return d;

def read_scalars(d):
    for quantity in d['dqs']:
        d[quantity]=[ d['x'+quantity].unpack_float() for k in range(d['nAll'])];
        del d['x'+quantity];
    del d['dqs'],d['nAll'];
    return d;

class LspOutput(file):
    '''represents an lsp output file on call,
       reads the header on open'''
    def __init__(self,filename,verbose=False,prefix='',buffering=-1):
        file.__init__(self,filename,"rb",buffering);
        self._get_header();
        self.verbose = verbose;
        self.prefix = prefix;
    def logprint(self,s):
        if self.verbose:
            if self.prefix == '':
                print('{}'.format(s));
            else:
                print('{}: {}'.format(self.prefix,s));
    def get_chunk(self, n):
        return self.read(n);
    def get_int(self):
        return xdr.Unpacker(self.read(4)).unpack_int();
    def get_uint(self):
        return xdr.Unpacker(self.read(4)).unpack_uint();
    def get_float(self):
        return xdr.Unpacker(self.read(4)).unpack_float();
    def get_str(self):
        l1 = self.get_int();
        l2 = self.get_int();
        if l1 != l2:
            print("warning, string prefixes are not equal...");
        size=l1;
        if l1%4:
            size+=4-l1%4;
        return xdr.Unpacker(self.get_chunk(size)).unpack_fstring(l1);
    def get_list(self,fmt):
        '''makes a list out of the fmt from the LspOutput f using the format
        i for int
        f for float
        d for double
        s for string'''
        out=[]
        for i in fmt:
            if i == 'i':
                out.append(self.get_int());
            elif i == 'f' or i == 'd':
                out.append(self.get_float());
            elif i == 's':
                out.append(self.get_str());
            else:
                raise ValueError("Unexpected flag '{}'".format(i));
        return out;

    def get_dict(self,fmt,keys):
        l=self.get_list(fmt);
        return dict(zip(keys,l));

    def _get_header(self):
        '''gets the header for the .p4 file'''
        self.header = self.get_dict('iiss',
                                    ['dump_type','dversion',
                                     'title','revision']);
        if self.header['dump_type'] == 2 or self.header['dump_type'] == 3:
            #this is a fields file or a scalars file.
            d = self.get_dict('fii',['timestamp','geometry','domains']);
            self.header.update(d);
            #reading quantities
            n = self.get_int();
            names=[self.get_str() for i in range(n)];
            units=[self.get_str() for i in range(n)];
            self.header['quantities'] = zip(names,units);
        elif self.header['dump_type'] == 6:
            #this is a particle movie file
            d = self.get_dict('iiii',
                              ['geometry','sflagsx','sflagsy','sflagsz']);
            self.header.update(d);
            #reading params
            n = self.get_int();
            flags=[bool(self.get_int()) for i in range(n)];
            units=[self.get_str() for i in range(n)];
            labels=['q','x','y','z','ux','uy','uz','E']
            if n == 8:
                pass;
            elif n == 7:
                labels = labels[:-1];
            elif n == 11:
                labels+=['xi','yi','zi'];
            else:
                raise ValueError('Incorrect number of parameters: {}'.format(n));
            self.header['params'] = [(i[0],i[1]) for i in zip(labels,units,flags)
                                     if i[2]];
        elif self.header['dump_type'] == 10:
            #this is a particle extraction file:
            self.header['geometry'] = self.get_int();
            #reading quantities
            n = self.get_int();
            self.header['quantities'] = [self.get_str() for i in range(n)];
        else:
            raise ValueError('Unknown dump_type: {}'.format(self.header['dump_type']));
        return;
    ###################
    #data processing
    def _out_getfields(self, var, pool_size,vector=True):
        if vector:
            size=3; call=read_fields;
            readin = set();
            for i in var:#we assume none of the fields end in x
                if i[-1] == 'x' or i[-1] == 'y' or i[-1] == 'z':
                    readin.add(i[:-1]);
                else:
                    readin.add(i);
        else:
            size=1; call=read_scalars;
            readin = set(var);
        doms = [];
        qs = [i[0] for i in self.header['quantities']];
        self.logprint('reading positions and making buffers');
        for i in range(self.header['domains']):
            self.logprint('reading domain {}'.format(i));
            iR, jR, kR = self.get_int(),self.get_int(),self.get_int();
            #getting grid parameters (real coordinates)
            nI = self.get_int(); Ip = [self.get_float() for i in range(nI)];
            nJ = self.get_int(); Jp = [self.get_float() for i in range(nJ)];
            nK = self.get_int(); Kp = [self.get_float() for i in range(nK)];
            self.logprint((nI,nJ,nK));
            nAll = nI*nJ*nK;
            d={}
            dqs=[];
            for quantity in qs:
                if quantity not in readin:
                    self.seek(nAll*4*size,1);
                else:
                    d.update({'x'+quantity:xdr.Unpacker(self.read(nAll*4*size))});
                    dqs.append(quantity);
            d.update({'nAll':nAll,'dqs': dqs,
                      'xp':Ip,'yp':Jp,'zp':Kp});
            doms.append(d);
        self.logprint('making points');
        pool=multiprocessing.Pool(pool_size);
        #making points
        doms[:] = pool.map(make_points,doms);
        self.logprint('converting buffers');
        doms[:] = pool.map(call,doms);
        pool.close();
        self.logprint('done! stringing together');
        for dom in doms[1:]:
            for k in doms[0]:
                doms[0][k].extend(dom[k]);
            del dom;
        return doms[0];

    def _getfields(self, var, vector=True):
        if vector:
            size=3;
            readin = set();
            for i in var:#we assume none of the fields end in x
                if i[-1] == 'x' or i[-1] == 'y' or i[-1] == 'z':
                    readin.add(i[:-1]);
                else:
                    readin.add(i);
        else:
            size=1; call=read_scalars;
            readin = set(var);
        doms = [];
        qs = [i[0] for i in self.header['quantities']];
        self.logprint('Reading positions and making buffers');
        for i in range(self.header['domains']):
            self.logprint('Reading domain {}.'.format(i));
            iR, jR, kR = self.get_int(),self.get_int(),self.get_int();
            #getting grid parameters (real coordinates)
            nI = self.get_int(); Ip = np.fromfile(self,dtype='>f4',count=nI);
            nJ = self.get_int(); Jp = np.fromfile(self,dtype='>f4',count=nJ);
            nK = self.get_int(); Kp = np.fromfile(self,dtype='>f4',count=nK);
            nAll = nI*nJ*nK;
            self.logprint('Dimensions are {}x{}x{}={}.'.format(nI,nJ,nK,nAll));
            d={}
            self.logprint('Making points.');
            #the way meshgrid works, it has to be in this order.
            d['y'], d['z'], d['x'] = np.vstack(np.meshgrid(Jp,Kp,Ip)).reshape(3,-1);
            for quantity in qs:
                if quantity not in readin:
                    self.seek(nAll*4*size,1);
                else:
                    self.logprint('Reading in {}'.format(quantity));
                    d[quantity] = np.fromfile(self,dtype='>f4',count=nAll*size);
                    if size==3:
                        data=d[quantity].reshape(nAll,size).T;
                        d[quantity+'x'],d[quantity+'y'],d[quantity+'z']= data;
                        del data;
            doms.append(d);
        self.logprint('Done! Stringing together.');
        out = { k : np.concatenate([i[k] for i in doms]) for k in doms[0] };
        return out;

    def _getmovie(self):
        params,_  = zip(*self.header['params']);
        nparams = len(params);
        pbytes = (nparams+1)*4;
        frames=[];
        while True:
            c=self.tell();
            self.read(1); #python, y u no eof?
            if self.tell() == c:
                break;
            self.seek(c);
            d=self.get_dict('fii',['t','step','pnum']);
            self.logprint('scanning frame at lsp step {}'.format(d['step']));
            d['pos']=self.tell();
            self.seek(d['pnum']*pbytes,1);
            frames.append(d);
        self.logprint('converting frames');
        for i,d in enumerate(frames):
            N = d['pnum'];
            lt=[('ip','>i4')]+zip(params,['>f4']*nparams);
            dt=np.dtype(lt);
            self.seek(d['pos']);
            arr=np.fromfile(self,dtype=dt,count=N);
            d['data']=arr;
            del d['pos'];
            self.logprint('done!');
            #checking eof
            c=self.tell();
            self.read(1); #python, y u no eof?
            self.logprint("Do we have eof? {}".format(self.tell() == c));
            frames[i] = d;
        return frames;

    def _getpext(self):
        nparams = len(self.header['quantities']);
        params = ['t','q','x','y','z','ux','uy','uz'];
        if nparams == 9:
            params+=['E'];
        elif nparams == 11:
            params+=['xi','yi','zi'];
        elif nparams == 12:
            params+=['E','xi','yi','zi'];
        #it's just floats here on out
        dt = zip(params, ['>f4']*len(params));
        out = np.fromfile(self,dtype=dt,count=-1);
        out = {k:out[k] for k in params};
        return out;

    def get_data(self,var=None):
        if not var and (self.header['dump_type']== 2 or self.header['dump_type']== 3):
            var=[i[0] for i in self.header['quantities']];
        if self.header['dump_type'] == 2:
            return self._getfields(var,vector=True);
        elif self.header['dump_type'] == 3:
            return self._getfields(var,vector=False);
        elif self.header['dump_type'] == 6:
            return self._getmovie();
        elif self.header['dump_type'] == 10:
            return self._getpext();
        else:
            return None;
    pass;
<commit_msg>minor edits<commit_after>'''
Reader for LSP output xdr files (.p4's)

'''
import xdrlib as xdr;
import itertools as itools;
import multiprocessing;
import numpy as np;

def make_points(d):
    d['x']=[]; d['y']=[]; d['z']=[];
    tmp1 = len(d['xp']);
    tmp2 = tmp1 * len(d['yp']);
    for i in range(d['nAll']):
        d['z'].append(d['zp'][i/tmp2]);
        d['y'].append(d['yp'][i%tmp2/tmp1]);
        d['x'].append(d['xp'][i%tmp1]);
    del d['zp'],d['yp'],d['xp'];
    return d;

class LspOutput(file):
    '''represents an lsp output file on call,
       reads the header on open'''
    def __init__(self,filename,verbose=False,prefix='',buffering=-1):
        file.__init__(self,filename,"rb",buffering);
        self._get_header();
        self.verbose = verbose;
        self.prefix = prefix;
    def logprint(self,s):
        if self.verbose:
            if self.prefix == '':
                print('{}'.format(s));
            else:
                print('{}: {}'.format(self.prefix,s));
    def get_chunk(self, n):
        return self.read(n);
    def get_int(self):
        return xdr.Unpacker(self.read(4)).unpack_int();
    def get_uint(self):
        return xdr.Unpacker(self.read(4)).unpack_uint();
    def get_float(self):
        return xdr.Unpacker(self.read(4)).unpack_float();
    def get_str(self):
        l1 = self.get_int();
        l2 = self.get_int();
        if l1 != l2:
            print("warning, string prefixes are not equal...");
        size=l1;
        if l1%4:
            size+=4-l1%4;
        return xdr.Unpacker(self.get_chunk(size)).unpack_fstring(l1);
    def get_list(self,fmt):
        '''makes a list out of the fmt from the LspOutput f using the format
        i for int
        f for float
        d for double
        s for string'''
        out=[]
        for i in fmt:
            if i == 'i':
                out.append(self.get_int());
            elif i == 'f' or i == 'd':
                out.append(self.get_float());
            elif i == 's':
                out.append(self.get_str());
            else:
                raise ValueError("Unexpected flag '{}'".format(i));
        return out;

    def get_dict(self,fmt,keys):
        l=self.get_list(fmt);
        return dict(zip(keys,l));

    def _get_header(self):
        '''gets the header for the .p4 file'''
        self.header = self.get_dict('iiss',
                                    ['dump_type','dversion',
                                     'title','revision']);
        if self.header['dump_type'] == 2 or self.header['dump_type'] == 3:
            #this is a fields file or a scalars file.
            d = self.get_dict('fii',['timestamp','geometry','domains']);
            self.header.update(d);
            #reading quantities
            n = self.get_int();
            names=[self.get_str() for i in range(n)];
            units=[self.get_str() for i in range(n)];
            self.header['quantities'] = zip(names,units);
        elif self.header['dump_type'] == 6:
            #this is a particle movie file
            d = self.get_dict('iiii',
                              ['geometry','sflagsx','sflagsy','sflagsz']);
            self.header.update(d);
            #reading params
            n = self.get_int();
            flags=[bool(self.get_int()) for i in range(n)];
            units=[self.get_str() for i in range(n)];
            labels=['q','x','y','z','ux','uy','uz','E']
            if n == 8:
                pass;
            elif n == 7:
                labels = labels[:-1];
            elif n == 11:
                labels+=['xi','yi','zi'];
            else:
                raise ValueError('Incorrect number of parameters: {}'.format(n));
            self.header['params'] = [(i[0],i[1]) for i in zip(labels,units,flags)
                                     if i[2]];
        elif self.header['dump_type'] == 10:
            #this is a particle extraction file:
            self.header['geometry'] = self.get_int();
            #reading quantities
            n = self.get_int();
            self.header['quantities'] = [self.get_str() for i in range(n)];
        else:
            raise ValueError('Unknown dump_type: {}'.format(self.header['dump_type']));
        return;
    ###################
    #data processing
    def _getfields(self, var, vector=True):
        if vector:
            size=3;
            readin = set();
            for i in var:#we assume none of the fields end in x
                if i[-1] == 'x' or i[-1] == 'y' or i[-1] == 'z':
                    readin.add(i[:-1]);
                else:
                    readin.add(i);
        else:
            size=1; call=read_scalars;
            readin = set(var);
        doms = [];
        qs = [i[0] for i in self.header['quantities']];
        self.logprint('Reading positions and making buffers');
        for i in range(self.header['domains']):
            self.logprint('Reading domain {}.'.format(i));
            iR, jR, kR = self.get_int(),self.get_int(),self.get_int();
            #getting grid parameters (real coordinates)
            nI = self.get_int(); Ip = np.fromfile(self,dtype='>f4',count=nI);
            nJ = self.get_int(); Jp = np.fromfile(self,dtype='>f4',count=nJ);
            nK = self.get_int(); Kp = np.fromfile(self,dtype='>f4',count=nK);
            nAll = nI*nJ*nK;
            self.logprint('Dimensions are {}x{}x{}={}.'.format(nI,nJ,nK,nAll));
            d={}
            self.logprint('Making points.');
            #the way meshgrid works, it has to be in this order.
            d['y'], d['z'], d['x'] = np.vstack(np.meshgrid(Jp,Kp,Ip)).reshape(3,-1);
            for quantity in qs:
                if quantity not in readin:
                    self.seek(nAll*4*size,1);
                else:
                    self.logprint('Reading in {}'.format(quantity));
                    d[quantity] = np.fromfile(self,dtype='>f4',count=nAll*size);
                    if size==3:
                        data=d[quantity].reshape(nAll,size).T;
                        d[quantity+'x'],d[quantity+'y'],d[quantity+'z']= data;
                        del data;
            doms.append(d);
        self.logprint('Done! Stringing together.');
        out = { k : np.concatenate([i[k] for i in doms]) for k in doms[0] };
        return out;

    def _getmovie(self):
        params,_  = zip(*self.header['params']);
        nparams = len(params);
        pbytes = (nparams+1)*4;
        frames=[];
        while True:
            c=self.tell();
            self.read(1); #python, y u no eof?
            if self.tell() == c:
                break;
            self.seek(c);
            d=self.get_dict('fii',['t','step','pnum']);
            self.logprint('scanning frame at lsp step {}'.format(d['step']));
            d['pos']=self.tell();
            self.seek(d['pnum']*pbytes,1);
            frames.append(d);
        self.logprint('converting frames');
        for i,d in enumerate(frames):
            N = d['pnum'];
            lt=[('ip','>i4')]+zip(params,['>f4']*nparams);
            dt=np.dtype(lt);
            self.seek(d['pos']);
            arr=np.fromfile(self,dtype=dt,count=N);
            d['data']=arr;
            del d['pos'];
            self.logprint('done!');
            #checking eof
            c=self.tell();
            self.read(1); #python, y u no eof?
            self.logprint("Do we have eof? {}".format(self.tell() == c));
            frames[i] = d;
        return frames;

    def _getpext(self):
        nparams = len(self.header['quantities']);
        params = ['t','q','x','y','z','ux','uy','uz'];
        if nparams == 9:
            params+=['E'];
        elif nparams == 11:
            params+=['xi','yi','zi'];
        elif nparams == 12:
            params+=['E','xi','yi','zi'];
        #it's just floats here on out
        dt = zip(params, ['>f4']*len(params));
        out = np.fromfile(self,dtype=dt,count=-1);
        out = {k:out[k] for k in params};
        return out;

    def get_data(self,var=None):
        if not var and (self.header['dump_type']== 2 or self.header['dump_type']== 3):
            var=[i[0] for i in self.header['quantities']];
        if self.header['dump_type'] == 2:
            return self._getfields(var,vector=True);
        elif self.header['dump_type'] == 3:
            return self._getfields(var,vector=False);
        elif self.header['dump_type'] == 6:
            return self._getmovie();
        elif self.header['dump_type'] == 10:
            return self._getpext();
        else:
            return None;
    pass;
<|endoftext|>


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

"""
Copyright 2016, Adam Kempenich. Questions? Comments? Feedback of any kind? Find me on GITHUB, @AdamKempenich
This utility was designed for use with devices compatible with the MagicHome Wifi app.

It currently supports:
   Bulbs (Firmware v.4 and greater)
   Legacy Bulbs (Firmware v.3 and lower)
   RGB Controllers
   RGB+WW Controllers
   RGB+WW+CW Controllers

##### Commands:
   Turning devices on and off
   Getting current device states
   Setting either colors, or whites (in bulbs)
   Setting colors (+ WW + CW) in compatible devices
   Sending preset functions

##### Tasks:
   Initial device setup via UDP.
   Custom commands (I know how it works ... just not a priority to implement currently.)
   Setting the device's internal clock.

##### Notes:
   Device types are:
      0: RGB
      1: RGB+WW
      2: RGB+WW+CW
      3: Bulb (v.4+)
      4: Bulb (v.3-)
      (Higher numbers reserved for future use)

   Functions accept integers 0-255 as parameters, or their equivalent byte format.

"""

import socket, csv, struct, datetime

class MagicHome_Wifi_Api:
   def __init__(self, device_ip, device_type, keep_alive=True):
      self.device_ip = device_ip
      self.device_type = device_type
      self.API_PORT = 5577
      self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      self.s.connect((self.device_ip, self.API_PORT))
      self.latest_connection = datetime.datetime.now()
      self.keep_alive = keep_alive

   def On(self):
      #Turns a device on
      self.Send_Bytes(0x71, 0x23, 0x0F, 0xA3) if self.device_type < 4 else self.Send_Bytes(0xCC, 0x23, 0x33)

   def Off(self):
      #Turns a device off
      self.Send_Bytes(device_ip, 0x71, 0x24, 0x0F, 0xA4) if self.device_type < 4 else self.Send_Bytes(0xCC, 0x24, 0x33)

   def Status(self):
      # Gets the current status of a device
      data = s.recv(14)
      self.Send_Bytes(device_ip, 0x81, 0x8A, 0x8B, 0x96)

   def Update_Device(self, r=0, g=0, b=0, white1=None, white2=None):
      # Updates a device based upon what we're sending to it
      # Values are excepted as integers between 0-255.
      # Whites can have a value of None.

      if self.device_type <= 1:
         # Update an RGB or an RGB + WW device
         white1 = self.Check_Number_Range(white1)
         message = [0x31, r, g, b, white1, 0x00, 0x0f]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 2:
         # Update an RGB + WW + CW device
         message = [0x31, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), self.Check_Number_Range(white1), self.Check_Number_Range(white2), 0x0f, 0x0f]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 3:
         # Update the white, or color, of a bulb
         if white1 != None:
            message = [0x31, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0x0f]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
         else:
            message = [0x31, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), 0x00, 0xf0, 0x0f]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 4:
         # Update the white, or color, of a legacy bulb
         if white1 != None:
            message = [0x56, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0xaa, 0x56, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0xaa]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
         else:
            message = [0x56, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), 0x00, 0xf0, 0xaa]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
      else:
         # Incompatible device received
         print "Incompatible device type received..."

   def Check_Number_Range(self, number):
      if number < 0:
         return 0
      elif number > 255:
         return 255
      else:
         return number

   def Send_Preset_Function(self, preset_number, speed):
      # Sends a preset command to a device
      # Presets can range from 0x25 (int 37) to 0x38 (int 56)
      if type <= 4:
         self.Send_Bytes(0xBB, preset_number, speed, 0x44)
      else:
         message = [0x61, preset_number, speed, 0x0F]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

   def Calculate_Checksum(self, bytes):
      # Calculates the checksum from an array of bytes
      return sum(bytes) & 0xFF

   def Send_Bytes(self, *bytes):
      # Sends commands to the device
      # If the device hasn't been communicated to in 5 minutes, reestablish the connection
      check_connection_time = abs((self.latest_connection-datetime.datetime.now()).total_seconds())

      if check_connection_time >= 290:
         self.s.connect((self.device_ip, self.API_PORT))
      message_length = len(bytes)
      self.s.send(struct.pack("B"*message_length, *bytes))
      # Close the connection unless requested not to
      if self.keep_alive == False:
         self.s.close
<commit_msg>Update magichome.py<commit_after>#!/usr/bin/env python

"""
Copyright 2016, Adam Kempenich. Questions? Comments? Feedback of any kind? Find me on GITHUB, @AdamKempenich
This utility was designed for use with devices compatible with the MagicHome Wifi app.

It currently supports:
   Bulbs (Firmware v.4 and greater)
   Legacy Bulbs (Firmware v.3 and lower)
   RGB Controllers
   RGB+WW Controllers
   RGB+WW+CW Controllers

##### Commands:
   Turning devices on and off
   Getting current device states
   Setting either colors, or whites (in bulbs)
   Setting colors (+ WW + CW) in compatible devices
   Sending preset functions

##### Tasks:
   Initial device setup via UDP.
   Custom commands (I know how it works ... just not a priority to implement currently.)
   Setting the device's internal clock.

##### Notes:
   Device types are:
      0: RGB
      1: RGB+WW
      2: RGB+WW+CW
      3: Bulb (v.4+)
      4: Bulb (v.3-)
      (Higher numbers reserved for future use)

   Functions accept integers 0-255 as parameters, or their equivalent byte format.

"""

import socket, csv, struct, datetime

class MagicHome_Wifi_Api:
   def __init__(self, device_ip, device_type, keep_alive=True):
      self.device_ip = device_ip
      self.device_type = device_type
      self.API_PORT = 5577
      self.latest_connection = datetime.datetime.now()
      self.keep_alive = keep_alive
      self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      self.s.settimeout(3)
      try:
         print "Trying first coneection"
         self.s.connect((self.device_ip, self.API_PORT))
      except socket.error, exc:
         print "Caught exception socket.error : %s" % exc
         if self.s:
             self.s.close()

   def On(self):
      #Turns a device on
      self.Send_Bytes(0x71, 0x23, 0x0F, 0xA3) if self.device_type < 4 else self.Send_Bytes(0xCC, 0x23, 0x33)

   def Off(self):
      #Turns a device off
      self.Send_Bytes(device_ip, 0x71, 0x24, 0x0F, 0xA4) if self.device_type < 4 else self.Send_Bytes(0xCC, 0x24, 0x33)

   def Status(self):
      # Gets the current status of a device
      data = s.recv(14)
      self.Send_Bytes(device_ip, 0x81, 0x8A, 0x8B, 0x96)

   def Update_Device(self, r=0, g=0, b=0, white1=None, white2=None):
      # Updates a device based upon what we're sending to it
      # Values are excepted as integers between 0-255.
      # Whites can have a value of None.

      if self.device_type <= 1:
         # Update an RGB or an RGB + WW device
         white1 = self.Check_Number_Range(white1)
         message = [0x31, r, g, b, white1, 0x00, 0x0f]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 2:
         # Update an RGB + WW + CW device
         message = [0x31, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), self.Check_Number_Range(white1), self.Check_Number_Range(white2), 0x0f, 0x0f]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 3:
         # Update the white, or color, of a bulb
         if white1 != None:
            message = [0x31, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0x0f]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
         else:
            message = [0x31, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), 0x00, 0xf0, 0x0f]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

      elif self.device_type == 4:
         # Update the white, or color, of a legacy bulb
         if white1 != None:
            message = [0x56, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0xaa, 0x56, 0x00, 0x00, 0x00, self.Check_Number_Range(white1), 0x0f, 0xaa]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
         else:
            message = [0x56, self.Check_Number_Range(r), self.Check_Number_Range(g), self.Check_Number_Range(b), 0x00, 0xf0, 0xaa]
            self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))
      else:
         # Incompatible device received
         print "Incompatible device type received..."

   def Check_Number_Range(self, number):
      if number < 0:
         return 0
      elif number > 255:
         return 255
      else:
         return number

   def Send_Preset_Function(self, preset_number, speed):
      # Sends a preset command to a device
      # Presets can range from 0x25 (int 37) to 0x38 (int 56)
      if type <= 4:
         self.Send_Bytes(0xBB, preset_number, speed, 0x44)
      else:
         message = [0x61, preset_number, speed, 0x0F]
         self.Send_Bytes(*(message+[self.Calculate_Checksum(message)]))

   def Calculate_Checksum(self, bytes):
      # Calculates the checksum from an array of bytes
      return sum(bytes) & 0xFF

   def Send_Bytes(self, *bytes):
      # Sends commands to the device
      # If the device hasn't been communicated to in 5 minutes, reestablish the connection
      check_connection_time = abs((self.latest_connection-datetime.datetime.now()).total_seconds())
      try:
         print "Trying second coneection"
         if check_connection_time >= 290:
            self.s.connect((self.device_ip, self.API_PORT))
         message_length = len(bytes)
         self.s.send(struct.pack("B"*message_length, *bytes))
         # Close the connection unless requested not to
         if self.keep_alive == False:
            self.s.close
      except socket.error, exc:
         print "Caught exception socket.error : %s" % exc
         if self.s:
             self.s.close()
<|endoftext|>


<commit_before># coding: utf-8

import functools
import re

from flask.ext import login
from flask.ext import oauth
from google.appengine.api import users
from google.appengine.ext import ndb
import flask
import unidecode

import config
import model
import task
import util

from main import app

_signals = flask.signals.Namespace()

###############################################################################
# Flask Login
###############################################################################
login_manager = login.LoginManager()

class AnonymousUser(login.AnonymousUserMixin):
  id = 0
  admin = False
  name = 'Anonymous'
  user_db = None

  def key(self):
    return None

  def has_permission(self, permission):
    return False

login_manager.anonymous_user = AnonymousUser

class FlaskUser(AnonymousUser):
  def __init__(self, user_db):
    self.user_db = user_db
    self.id = user_db.key.id()
    self.name = user_db.name
    self.admin = user_db.admin

  def key(self):
    return self.user_db.key.urlsafe()

  def get_id(self):
    return self.user_db.key.urlsafe()

  def is_authenticated(self):
    return True

  def is_active(self):
    return self.user_db.active

  def is_anonymous(self):
    return False

  def has_permission(self, permission):
    return self.user_db.has_permission(permission)

@login_manager.user_loader
def load_user(key):
  user_db = ndb.Key(urlsafe=key).get()
  if user_db:
    return FlaskUser(user_db)
  return None

login_manager.init_app(app)

def current_user_id():
  return login.current_user.id

def current_user_key():
  return login.current_user.user_db.key if login.current_user.user_db else None

def current_user_db():
  return login.current_user.user_db

def is_logged_in():
  return login.current_user.id != 0

###############################################################################
# Decorators
###############################################################################
def login_required(f):
  decorator_order_guard(f, 'auth.login_required')

  @functools.wraps(f)
  def decorated_function(*args, **kws):
    if is_logged_in():
      return f(*args, **kws)
    if flask.request.path.startswith('/_s/'):
      return flask.abort(401)
    return flask.redirect(flask.url_for('signin', next=flask.request.url))
  return decorated_function

def admin_required(f):
  decorator_order_guard(f, 'auth.admin_required')

  @functools.wraps(f)
  def decorated_function(*args, **kws):
    if is_logged_in() and current_user_db().admin:
      return f(*args, **kws)
    if not is_logged_in() and flask.request.path.startswith('/_s/'):
      return flask.abort(401)
    if not is_logged_in():
      return flask.redirect(flask.url_for('signin', next=flask.request.url))
    return flask.abort(403)
  return decorated_function

permission_registered = _signals.signal('permission-registered')

def permission_required(permission=None, methods=None):
  def permission_decorator(f):
    decorator_order_guard(f, 'auth.permission_required')

    # default to decorated function name as permission
    perm = permission or f.func_name
    meths = [m.upper() for m in methods] if methods else None

    permission_registered.send(f, permission=perm)

    @functools.wraps(f)
    def decorated_function(*args, **kws):
      if meths and flask.request.method.upper() not in meths:
        return f(*args, **kws)
      if is_logged_in() and current_user_db().has_permission(perm):
        return f(*args, **kws)
      if not is_logged_in():
        if flask.request.path.startswith('/_s/'):
          return flask.abort(401)
        return flask.redirect(flask.url_for('signin', next=flask.request.url))
      return flask.abort(403)
    return decorated_function
  return permission_decorator

###############################################################################
# Sign in stuff
###############################################################################
@app.route('/login/')
@app.route('/signin/')
def signin():
  next_url = util.get_next_url()

  google_signin_url = flask.url_for('signin_google', next=next_url)
  twitter_signin_url = flask.url_for('signin_twitter', next=next_url)
  facebook_signin_url = flask.url_for('signin_facebook', next=next_url)

  return flask.render_template(
      'signin.html',
      title='Please sign in',
      html_class='signin',
      google_signin_url=google_signin_url,
      twitter_signin_url=twitter_signin_url,
      facebook_signin_url=facebook_signin_url,
      next_url=next_url,
    )

@app.route('/signout/')
def signout():
  login.logout_user()
  flask.flash(u'You have been signed out.', category='success')
  return flask.redirect(flask.url_for('welcome'))

###############################################################################
# Google
###############################################################################
@app.route('/signin/google/')
def signin_google():
  save_request_params()
  google_url = users.create_login_url(flask.url_for('google_authorized'))
  return flask.redirect(google_url)

@app.route('/_s/callback/google/authorized/')
def google_authorized():
  google_user = users.get_current_user()
  if google_user is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  user_db = retrieve_user_from_google(google_user)
  return signin_user_db(user_db)

def retrieve_user_from_google(google_user):
  auth_id = 'federated_%s' % google_user.user_id()
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    if not user_db.admin and users.is_current_user_admin():
      user_db.admin = True
      user_db.put()
    return user_db

  return create_user_db(
      auth_id,
      re.sub(r'_+|-+|\.+', ' ', google_user.email().split('@')[0]).title(),
      google_user.email(),
      google_user.email(),
      verified=True,
      admin=users.is_current_user_admin(),
    )

###############################################################################
# Twitter
###############################################################################
twitter_oauth = oauth.OAuth()

twitter = twitter_oauth.remote_app(
    'twitter',
    base_url='https://api.twitter.com/1.1/',
    request_token_url='https://api.twitter.com/oauth/request_token',
    access_token_url='https://api.twitter.com/oauth/access_token',
    authorize_url='https://api.twitter.com/oauth/authorize',
    consumer_key=config.CONFIG_DB.twitter_consumer_key,
    consumer_secret=config.CONFIG_DB.twitter_consumer_secret,
  )

@app.route('/_s/callback/twitter/oauth-authorized/')
@twitter.authorized_handler
def twitter_authorized(resp):
  if resp is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  flask.session['oauth_token'] = (
      resp['oauth_token'],
      resp['oauth_token_secret'],
    )
  user_db = retrieve_user_from_twitter(resp)
  return signin_user_db(user_db)

@twitter.tokengetter
def get_twitter_token():
  return flask.session.get('oauth_token')

@app.route('/signin/twitter/')
def signin_twitter():
  flask.session.pop('oauth_token', None)
  save_request_params()
  try:
    return twitter.authorize(callback=flask.url_for('twitter_authorized'))
  except:
    flask.flash(
        'Something went wrong with Twitter sign in. Please try again.',
        category='danger',
      )
    return flask.redirect(flask.url_for('signin', next=util.get_next_url()))

def retrieve_user_from_twitter(response):
  auth_id = 'twitter_%s' % response['user_id']
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    return user_db

  return create_user_db(
      auth_id,
      response['screen_name'],
      response['screen_name'],
    )

###############################################################################
# Facebook
###############################################################################
facebook_oauth = oauth.OAuth()

facebook = facebook_oauth.remote_app(
    'facebook',
    base_url='https://graph.facebook.com/',
    request_token_url=None,
    access_token_url='/oauth/access_token',
    authorize_url='https://www.facebook.com/dialog/oauth',
    consumer_key=config.CONFIG_DB.facebook_app_id,
    consumer_secret=config.CONFIG_DB.facebook_app_secret,
    request_token_params={'scope': 'email'},
  )

@app.route('/_s/callback/facebook/oauth-authorized/')
@facebook.authorized_handler
def facebook_authorized(resp):
  if resp is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  flask.session['oauth_token'] = (resp['access_token'], '')
  me = facebook.get('/me')
  user_db = retrieve_user_from_facebook(me.data)
  return signin_user_db(user_db)

@facebook.tokengetter
def get_facebook_oauth_token():
  return flask.session.get('oauth_token')

@app.route('/signin/facebook/')
def signin_facebook():
  save_request_params()
  return facebook.authorize(callback=flask.url_for(
      'facebook_authorized', _external=True
    ))

def retrieve_user_from_facebook(response):
  auth_id = 'facebook_%s' % response['id']
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    return user_db
  return create_user_db(
      auth_id,
      response['name'],
      response.get('username', response['name']),
      response.get('email', ''),
      verified=bool(response.get('email', '')),
    )

###############################################################################
# Helpers
###############################################################################
def decorator_order_guard(f, decorator_name):
  if f in app.view_functions.values():
    raise SyntaxError(
        'Do not use %s above app.route decorators as it would not be checked. '
        'Instead move the line below the app.route lines.' % decorator_name
      )

def create_user_db(auth_id, name, username, email='', verified=False, **params):
  email = email.lower()
  if verified and email:
    user_dbs, _ = model.User.get_dbs(email=email, verified=True, limit=2)
    if len(user_dbs) == 1:
      user_db = user_dbs[0]
      user_db.auth_ids.append(auth_id)
      user_db.put()
      task.new_user_notification(user_db)
      return user_db

  if isinstance(username, str):
    username = username.decode('utf-8')
  username = unidecode.unidecode(username.split('@')[0].lower()).strip()
  username = re.sub(r'[\W_]+', '.', username)
  new_username = username
  n = 1
  while not model.User.is_username_available(new_username):
    new_username = '%s%d' % (username, n)
    n += 1

  user_db = model.User(
      name=name,
      email=email,
      username=new_username,
      auth_ids=[auth_id],
      verified=verified,
      token=util.uuid(),
      **params
    )
  user_db.put()
  task.new_user_notification(user_db)
  return user_db

def save_request_params():
  flask.session['auth-params'] = {
      'next': util.get_next_url(),
      'remember': util.param('remember', bool),
    }

@ndb.toplevel
def signin_user_db(user_db):
  if not user_db:
    return flask.redirect(flask.url_for('signin'))
  flask_user_db = FlaskUser(user_db)
  auth_params = flask.session.get('auth-params', {
      'next': flask.url_for('welcome'),
      'remember': False,
    })
  if login.login_user(flask_user_db, remember=auth_params['remember']):
    user_db.put_async()
    flask.flash('Hello %s, welcome to %s.' % (
        user_db.name, config.CONFIG_DB.brand_name,
      ), category='success')
    return flask.redirect(util.get_next_url(auth_params['next']))
  flask.flash('Sorry, but you could not sign in.', category='danger')
  return flask.redirect(flask.url_for('signin'))
<commit_msg>After signing out redirect to sign in<commit_after># coding: utf-8

import functools
import re

from flask.ext import login
from flask.ext import oauth
from google.appengine.api import users
from google.appengine.ext import ndb
import flask
import unidecode

import config
import model
import task
import util

from main import app

_signals = flask.signals.Namespace()

###############################################################################
# Flask Login
###############################################################################
login_manager = login.LoginManager()

class AnonymousUser(login.AnonymousUserMixin):
  id = 0
  admin = False
  name = 'Anonymous'
  user_db = None

  def key(self):
    return None

  def has_permission(self, permission):
    return False

login_manager.anonymous_user = AnonymousUser

class FlaskUser(AnonymousUser):
  def __init__(self, user_db):
    self.user_db = user_db
    self.id = user_db.key.id()
    self.name = user_db.name
    self.admin = user_db.admin

  def key(self):
    return self.user_db.key.urlsafe()

  def get_id(self):
    return self.user_db.key.urlsafe()

  def is_authenticated(self):
    return True

  def is_active(self):
    return self.user_db.active

  def is_anonymous(self):
    return False

  def has_permission(self, permission):
    return self.user_db.has_permission(permission)

@login_manager.user_loader
def load_user(key):
  user_db = ndb.Key(urlsafe=key).get()
  if user_db:
    return FlaskUser(user_db)
  return None

login_manager.init_app(app)

def current_user_id():
  return login.current_user.id

def current_user_key():
  return login.current_user.user_db.key if login.current_user.user_db else None

def current_user_db():
  return login.current_user.user_db

def is_logged_in():
  return login.current_user.id != 0

###############################################################################
# Decorators
###############################################################################
def login_required(f):
  decorator_order_guard(f, 'auth.login_required')

  @functools.wraps(f)
  def decorated_function(*args, **kws):
    if is_logged_in():
      return f(*args, **kws)
    if flask.request.path.startswith('/_s/'):
      return flask.abort(401)
    return flask.redirect(flask.url_for('signin', next=flask.request.url))
  return decorated_function

def admin_required(f):
  decorator_order_guard(f, 'auth.admin_required')

  @functools.wraps(f)
  def decorated_function(*args, **kws):
    if is_logged_in() and current_user_db().admin:
      return f(*args, **kws)
    if not is_logged_in() and flask.request.path.startswith('/_s/'):
      return flask.abort(401)
    if not is_logged_in():
      return flask.redirect(flask.url_for('signin', next=flask.request.url))
    return flask.abort(403)
  return decorated_function

permission_registered = _signals.signal('permission-registered')

def permission_required(permission=None, methods=None):
  def permission_decorator(f):
    decorator_order_guard(f, 'auth.permission_required')

    # default to decorated function name as permission
    perm = permission or f.func_name
    meths = [m.upper() for m in methods] if methods else None

    permission_registered.send(f, permission=perm)

    @functools.wraps(f)
    def decorated_function(*args, **kws):
      if meths and flask.request.method.upper() not in meths:
        return f(*args, **kws)
      if is_logged_in() and current_user_db().has_permission(perm):
        return f(*args, **kws)
      if not is_logged_in():
        if flask.request.path.startswith('/_s/'):
          return flask.abort(401)
        return flask.redirect(flask.url_for('signin', next=flask.request.url))
      return flask.abort(403)
    return decorated_function
  return permission_decorator

###############################################################################
# Sign in stuff
###############################################################################
@app.route('/login/')
@app.route('/signin/')
def signin():
  next_url = util.get_next_url()

  google_signin_url = flask.url_for('signin_google', next=next_url)
  twitter_signin_url = flask.url_for('signin_twitter', next=next_url)
  facebook_signin_url = flask.url_for('signin_facebook', next=next_url)

  return flask.render_template(
      'signin.html',
      title='Please sign in',
      html_class='signin',
      google_signin_url=google_signin_url,
      twitter_signin_url=twitter_signin_url,
      facebook_signin_url=facebook_signin_url,
      next_url=next_url,
    )

@app.route('/signout/')
def signout():
  login.logout_user()
  flask.flash(u'You have been signed out.', category='success')
  return flask.redirect(flask.url_for('signin'))

###############################################################################
# Google
###############################################################################
@app.route('/signin/google/')
def signin_google():
  save_request_params()
  google_url = users.create_login_url(flask.url_for('google_authorized'))
  return flask.redirect(google_url)

@app.route('/_s/callback/google/authorized/')
def google_authorized():
  google_user = users.get_current_user()
  if google_user is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  user_db = retrieve_user_from_google(google_user)
  return signin_user_db(user_db)

def retrieve_user_from_google(google_user):
  auth_id = 'federated_%s' % google_user.user_id()
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    if not user_db.admin and users.is_current_user_admin():
      user_db.admin = True
      user_db.put()
    return user_db

  return create_user_db(
      auth_id,
      re.sub(r'_+|-+|\.+', ' ', google_user.email().split('@')[0]).title(),
      google_user.email(),
      google_user.email(),
      verified=True,
      admin=users.is_current_user_admin(),
    )

###############################################################################
# Twitter
###############################################################################
twitter_oauth = oauth.OAuth()

twitter = twitter_oauth.remote_app(
    'twitter',
    base_url='https://api.twitter.com/1.1/',
    request_token_url='https://api.twitter.com/oauth/request_token',
    access_token_url='https://api.twitter.com/oauth/access_token',
    authorize_url='https://api.twitter.com/oauth/authorize',
    consumer_key=config.CONFIG_DB.twitter_consumer_key,
    consumer_secret=config.CONFIG_DB.twitter_consumer_secret,
  )

@app.route('/_s/callback/twitter/oauth-authorized/')
@twitter.authorized_handler
def twitter_authorized(resp):
  if resp is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  flask.session['oauth_token'] = (
      resp['oauth_token'],
      resp['oauth_token_secret'],
    )
  user_db = retrieve_user_from_twitter(resp)
  return signin_user_db(user_db)

@twitter.tokengetter
def get_twitter_token():
  return flask.session.get('oauth_token')

@app.route('/signin/twitter/')
def signin_twitter():
  flask.session.pop('oauth_token', None)
  save_request_params()
  try:
    return twitter.authorize(callback=flask.url_for('twitter_authorized'))
  except:
    flask.flash(
        'Something went wrong with Twitter sign in. Please try again.',
        category='danger',
      )
    return flask.redirect(flask.url_for('signin', next=util.get_next_url()))

def retrieve_user_from_twitter(response):
  auth_id = 'twitter_%s' % response['user_id']
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    return user_db

  return create_user_db(
      auth_id,
      response['screen_name'],
      response['screen_name'],
    )

###############################################################################
# Facebook
###############################################################################
facebook_oauth = oauth.OAuth()

facebook = facebook_oauth.remote_app(
    'facebook',
    base_url='https://graph.facebook.com/',
    request_token_url=None,
    access_token_url='/oauth/access_token',
    authorize_url='https://www.facebook.com/dialog/oauth',
    consumer_key=config.CONFIG_DB.facebook_app_id,
    consumer_secret=config.CONFIG_DB.facebook_app_secret,
    request_token_params={'scope': 'email'},
  )

@app.route('/_s/callback/facebook/oauth-authorized/')
@facebook.authorized_handler
def facebook_authorized(resp):
  if resp is None:
    flask.flash(u'You denied the request to sign in.')
    return flask.redirect(util.get_next_url())

  flask.session['oauth_token'] = (resp['access_token'], '')
  me = facebook.get('/me')
  user_db = retrieve_user_from_facebook(me.data)
  return signin_user_db(user_db)

@facebook.tokengetter
def get_facebook_oauth_token():
  return flask.session.get('oauth_token')

@app.route('/signin/facebook/')
def signin_facebook():
  save_request_params()
  return facebook.authorize(callback=flask.url_for(
      'facebook_authorized', _external=True
    ))

def retrieve_user_from_facebook(response):
  auth_id = 'facebook_%s' % response['id']
  user_db = model.User.get_by('auth_ids', auth_id)
  if user_db:
    return user_db
  return create_user_db(
      auth_id,
      response['name'],
      response.get('username', response['name']),
      response.get('email', ''),
      verified=bool(response.get('email', '')),
    )

###############################################################################
# Helpers
###############################################################################
def decorator_order_guard(f, decorator_name):
  if f in app.view_functions.values():
    raise SyntaxError(
        'Do not use %s above app.route decorators as it would not be checked. '
        'Instead move the line below the app.route lines.' % decorator_name
      )

def create_user_db(auth_id, name, username, email='', verified=False, **params):
  email = email.lower()
  if verified and email:
    user_dbs, _ = model.User.get_dbs(email=email, verified=True, limit=2)
    if len(user_dbs) == 1:
      user_db = user_dbs[0]
      user_db.auth_ids.append(auth_id)
      user_db.put()
      task.new_user_notification(user_db)
      return user_db

  if isinstance(username, str):
    username = username.decode('utf-8')
  username = unidecode.unidecode(username.split('@')[0].lower()).strip()
  username = re.sub(r'[\W_]+', '.', username)
  new_username = username
  n = 1
  while not model.User.is_username_available(new_username):
    new_username = '%s%d' % (username, n)
    n += 1

  user_db = model.User(
      name=name,
      email=email,
      username=new_username,
      auth_ids=[auth_id],
      verified=verified,
      token=util.uuid(),
      **params
    )
  user_db.put()
  task.new_user_notification(user_db)
  return user_db

def save_request_params():
  flask.session['auth-params'] = {
      'next': util.get_next_url(),
      'remember': util.param('remember', bool),
    }

@ndb.toplevel
def signin_user_db(user_db):
  if not user_db:
    return flask.redirect(flask.url_for('signin'))
  flask_user_db = FlaskUser(user_db)
  auth_params = flask.session.get('auth-params', {
      'next': flask.url_for('welcome'),
      'remember': False,
    })
  if login.login_user(flask_user_db, remember=auth_params['remember']):
    user_db.put_async()
    flask.flash('Hello %s, welcome to %s.' % (
        user_db.name, config.CONFIG_DB.brand_name,
      ), category='success')
    return flask.redirect(util.get_next_url(auth_params['next']))
  flask.flash('Sorry, but you could not sign in.', category='danger')
  return flask.redirect(flask.url_for('signin'))
<|endoftext|>


<commit_before>import sys
sys.path.insert(0, 'lib.zip')

import flask
from flaskext import wtf
import config

app = flask.Flask(__name__)
app.config.from_object(config)

from google.appengine.api import mail

import auth
import util
import model
import admin

@app.route('/')
def welcome():
  return flask.render_template(
      'welcome.html',
      html_class='welcome',
      channel_name='welcome',
    )

################################################################################
# Profile stuff
################################################################################
class ProfileUpdateForm(wtf.Form):
  name = wtf.TextField('Name', [wtf.validators.required()])
  email = wtf.TextField('Email', [
      wtf.validators.optional(),
      wtf.validators.email("That doesn't look like an email"),
    ])

@app.route('/_s/profile/', endpoint='profile_service')
@app.route('/profile/', methods=['GET', 'POST'], endpoint='profile')
@auth.login_required
def profile():
  form = ProfileUpdateForm()
  user_db = auth.current_user_db()
  if form.validate_on_submit():
    user_db.name = form.name.data
    user_db.email = form.email.data.lower()
    user_db.put()
    return flask.redirect(flask.url_for('welcome'))
  if not form.errors:
    form.name.data = user_db.name
    form.email.data = user_db.email or ''

  if flask.request.path.startswith('/_s/'):
    return util.jsonify_model_db(user_db)

  return flask.render_template(
      'profile.html',
      title='Profile',
      html_class='profile',
      form=form,
      user_db=user_db,
    )

################################################################################
# Feedback
################################################################################
class FeedbackForm(wtf.Form):
  subject = wtf.TextField('Subject', [wtf.validators.required()])
  message = wtf.TextAreaField('Message', [wtf.validators.required()])
  email = wtf.TextField('Email (optional)', [
      wtf.validators.optional(),
      wtf.validators.email("That doesn't look like an email"),
    ])

@app.route('/feedback/', methods=['GET', 'POST'])
def feedback():
  form = FeedbackForm()
  if form.validate_on_submit():
    mail.send_mail(
        sender=config.CONFIG_DB.feedback_email,
        to=config.CONFIG_DB.feedback_email,
        subject='[%s] %s' % (
            config.CONFIG_DB.brand_name,
            form.subject.data,
          ),
        reply_to=form.email.data or config.CONFIG_DB.feedback_email,
        body='%s\n\n%s' % (form.message.data, form.email.data)
      )
    flask.flash('Thank you for your feedback!', category='success')
    return flask.redirect(flask.url_for('welcome'))
  if not form.errors and auth.current_user_id() > 0:
    form.email.data = auth.current_user_db().email

  return flask.render_template(
      'feedback.html',
      title='Feedback',
      html_class='feedback',
      form=form,
    )

################################################################################
# User Stuff
################################################################################
@app.route('/_s/user/', endpoint='user_list_service')
@app.route('/user/', endpoint='user_list')
def user_list():
  user_dbs, more_cursor = util.retrieve_dbs(
      model.User,
      model.User.query(),
      limit=util.param('limit', int),
      cursor=util.param('cursor'),
      order=util.param('order'),
      name=util.param('name'),
    )

  if flask.request.path.startswith('/_s/'):
    return util.jsonify_model_dbs(user_dbs, more_cursor)

  return flask.render_template(
      'user_list.html',
      html_class='user',
      title='User List',
      user_dbs=user_dbs,
      more_url=util.generate_more_url(more_cursor),
    )

################################################################################
# Extras
################################################################################
@app.route('/_s/extras/', endpoint='extras_service')
@app.route('/extras/', endpoint='extras')
def extras():
  country = None
  region = None
  city = None
  city_lat_long = None
  if 'X-AppEngine-Country' in flask.request.headers:
    country = flask.request.headers['X-AppEngine-Country']
  if 'X-AppEngine-Region' in flask.request.headers:
    region = flask.request.headers['X-AppEngine-Region']
  if 'X-AppEngine-City' in flask.request.headers:
    city = flask.request.headers['X-AppEngine-City']
  if 'X-AppEngine-CityLatLong' in flask.request.headers:
    city_lat_long = flask.request.headers['X-AppEngine-CityLatLong']

  extra_info = {
    'country': country,
    'region': region,
    'city': city,
    'city_lat_long': city_lat_long,
    'user_agent': flask.request.headers['User-Agent'],
  }

  if flask.request.path.startswith('/_s/'):
    return flask.jsonify(extra_info)

  return flask.render_template(
      'extras.html',
      html_class='extras',
      title='Extras',
      extra_info=extra_info,
    )

@app.route('/chat/')
def chat():
  return flask.render_template(
      'chat.html',
      title='Chat',
      html_class='chat',
      channel_name='chat',
    )

################################################################################
# Error Handling
################################################################################
@app.errorhandler(400)
@app.errorhandler(401)
@app.errorhandler(403)
@app.errorhandler(404)
@app.errorhandler(410)
@app.errorhandler(418)
@app.errorhandler(500)
def error_handler(e):
  try:
    e.code
  except:
    class e(object):
      code = 500
      name = 'Internal Server Error'

  if flask.request.path.startswith('/_s/'):
    return flask.jsonify({
        'status': 'error',
        'error_code': e.code,
        'error_name': e.name.lower().replace(' ', '_'),
        'error_message': e.name,
      })

  return flask.render_template(
      'error.html',
      title='%s!!1' % (e.name),
      html_class='error-page',
      error=e,
    ), e.code
<commit_msg>Included the error code in the title for error handling pages<commit_after>import sys
sys.path.insert(0, 'lib.zip')

import flask
from flaskext import wtf
import config

app = flask.Flask(__name__)
app.config.from_object(config)

from google.appengine.api import mail

import auth
import util
import model
import admin

@app.route('/')
def welcome():
  return flask.render_template(
      'welcome.html',
      html_class='welcome',
      channel_name='welcome',
    )

################################################################################
# Profile stuff
################################################################################
class ProfileUpdateForm(wtf.Form):
  name = wtf.TextField('Name', [wtf.validators.required()])
  email = wtf.TextField('Email', [
      wtf.validators.optional(),
      wtf.validators.email("That doesn't look like an email"),
    ])

@app.route('/_s/profile/', endpoint='profile_service')
@app.route('/profile/', methods=['GET', 'POST'], endpoint='profile')
@auth.login_required
def profile():
  form = ProfileUpdateForm()
  user_db = auth.current_user_db()
  if form.validate_on_submit():
    user_db.name = form.name.data
    user_db.email = form.email.data.lower()
    user_db.put()
    return flask.redirect(flask.url_for('welcome'))
  if not form.errors:
    form.name.data = user_db.name
    form.email.data = user_db.email or ''

  if flask.request.path.startswith('/_s/'):
    return util.jsonify_model_db(user_db)

  return flask.render_template(
      'profile.html',
      title='Profile',
      html_class='profile',
      form=form,
      user_db=user_db,
    )

################################################################################
# Feedback
################################################################################
class FeedbackForm(wtf.Form):
  subject = wtf.TextField('Subject', [wtf.validators.required()])
  message = wtf.TextAreaField('Message', [wtf.validators.required()])
  email = wtf.TextField('Email (optional)', [
      wtf.validators.optional(),
      wtf.validators.email("That doesn't look like an email"),
    ])

@app.route('/feedback/', methods=['GET', 'POST'])
def feedback():
  form = FeedbackForm()
  if form.validate_on_submit():
    mail.send_mail(
        sender=config.CONFIG_DB.feedback_email,
        to=config.CONFIG_DB.feedback_email,
        subject='[%s] %s' % (
            config.CONFIG_DB.brand_name,
            form.subject.data,
          ),
        reply_to=form.email.data or config.CONFIG_DB.feedback_email,
        body='%s\n\n%s' % (form.message.data, form.email.data)
      )
    flask.flash('Thank you for your feedback!', category='success')
    return flask.redirect(flask.url_for('welcome'))
  if not form.errors and auth.current_user_id() > 0:
    form.email.data = auth.current_user_db().email

  return flask.render_template(
      'feedback.html',
      title='Feedback',
      html_class='feedback',
      form=form,
    )

################################################################################
# User Stuff
################################################################################
@app.route('/_s/user/', endpoint='user_list_service')
@app.route('/user/', endpoint='user_list')
def user_list():
  user_dbs, more_cursor = util.retrieve_dbs(
      model.User,
      model.User.query(),
      limit=util.param('limit', int),
      cursor=util.param('cursor'),
      order=util.param('order'),
      name=util.param('name'),
    )

  if flask.request.path.startswith('/_s/'):
    return util.jsonify_model_dbs(user_dbs, more_cursor)

  return flask.render_template(
      'user_list.html',
      html_class='user',
      title='User List',
      user_dbs=user_dbs,
      more_url=util.generate_more_url(more_cursor),
    )

################################################################################
# Extras
################################################################################
@app.route('/_s/extras/', endpoint='extras_service')
@app.route('/extras/', endpoint='extras')
def extras():
  country = None
  region = None
  city = None
  city_lat_long = None
  if 'X-AppEngine-Country' in flask.request.headers:
    country = flask.request.headers['X-AppEngine-Country']
  if 'X-AppEngine-Region' in flask.request.headers:
    region = flask.request.headers['X-AppEngine-Region']
  if 'X-AppEngine-City' in flask.request.headers:
    city = flask.request.headers['X-AppEngine-City']
  if 'X-AppEngine-CityLatLong' in flask.request.headers:
    city_lat_long = flask.request.headers['X-AppEngine-CityLatLong']

  extra_info = {
    'country': country,
    'region': region,
    'city': city,
    'city_lat_long': city_lat_long,
    'user_agent': flask.request.headers['User-Agent'],
  }

  if flask.request.path.startswith('/_s/'):
    return flask.jsonify(extra_info)

  return flask.render_template(
      'extras.html',
      html_class='extras',
      title='Extras',
      extra_info=extra_info,
    )

@app.route('/chat/')
def chat():
  return flask.render_template(
      'chat.html',
      title='Chat',
      html_class='chat',
      channel_name='chat',
    )

################################################################################
# Error Handling
################################################################################
@app.errorhandler(400)
@app.errorhandler(401)
@app.errorhandler(403)
@app.errorhandler(404)
@app.errorhandler(410)
@app.errorhandler(418)
@app.errorhandler(500)
def error_handler(e):
  try:
    e.code
  except:
    class e(object):
      code = 500
      name = 'Internal Server Error'

  if flask.request.path.startswith('/_s/'):
    return flask.jsonify({
        'status': 'error',
        'error_code': e.code,
        'error_name': e.name.lower().replace(' ', '_'),
        'error_message': e.name,
      })

  return flask.render_template(
      'error.html',
      title='Error %d (%s)!!1' % (e.code, e.name),
      html_class='error-page',
      error=e,
    ), e.code
<|endoftext|>