<commit_before>import wx
import wx.lib.newevent
import math
import cmath
import yaml

import pyggdrasil

TreeChangedEvent, TREE_CHANGED_EVENT = wx.lib.newevent.NewEvent()
GraphSelectedEvent, GRAPH_SELECTED_EVENT = wx.lib.newevent.NewEvent()

class Main(wx.Frame):
    def __init__(self, root=None, filename=None, *args, **kwargs):
        wx.Frame.__init__(self, *args, **kwargs)

        self.root = root or pyggdrasil.model.Node('root', None)
        self.filename = filename

        menubar = wx.MenuBar()

        file = wx.Menu()
        new = file.Append(wx.ID_NEW)
        self.Bind(wx.EVT_MENU, self.OnNew, new)
        open = file.Append(wx.ID_OPEN)
        self.Bind(wx.EVT_MENU, self.OnOpen, open)
        file.AppendSeparator()

        close = file.Append(wx.ID_CLOSE, 'Close\tCtrl-W')
        self.Bind(wx.EVT_MENU, self.OnClose, close)
        save = file.Append(wx.ID_SAVE)
        self.Bind(wx.EVT_MENU, self.OnSave, save)
        saveas = file.Append(wx.ID_SAVEAS)
        self.Bind(wx.EVT_MENU, self.OnSaveAs, saveas)
        file.AppendSeparator()

        file.Append(wx.ID_EXIT)
        menubar.Append(file, '&File')

        self.SetMenuBar(menubar)

        box = wx.BoxSizer(wx.HORIZONTAL)

        self._tree = Tree(self.root, self)
        box.Add(self._tree, 0, wx.EXPAND)

        self.graph = Graph(self.root, 40, 5, self)
        box.Add(self.graph, 1, wx.EXPAND)

        self.SetSizer(box)

        self._tree.Bind(TREE_CHANGED_EVENT, self.OnTreeChange)
        self._tree.Bind(wx.EVT_TREE_SEL_CHANGED, self.OnTreeSelected)
        self.graph.Bind(GRAPH_SELECTED_EVENT, self.OnGraphSelected)

    def getfilename(self):
        return self._filename
    def setfilename(self, value):
        self._filename = value
        self.SetTitle('Pyggdrasil - ' + (self.filename or 'new'))
    filename = property(getfilename, setfilename)

    def OnNew(self, event):
        frame = Main(parent=self.GetParent(), id=wx.ID_ANY)
        frame.Show(True)

    def OnOpen(self, event):
        filename = wx.LoadFileSelector('Pyggdrasil', extension='pyg', parent=self)
        if filename:
            file = open(filename)
            try:
                root = pyggdrasil.model.Node.from_raw(yaml.load(file))
                frame = Main(root, filename, self.GetParent(), wx.ID_ANY)
                frame.Show(True)
            finally:
                file.close()

    def OnSave(self, event):
        if self.filename:
            self._Save()
        else:
            self.OnSaveAs(event)

    def OnSaveAs(self, event):
        filename = wx.SaveFileSelector('Pyggdrasil', extension='pyg', parent=self)
        if filename:
            if '.' not in filename:
                filename += '.pyg'
            self.filename = filename
            self._Save()

    def _Save(self):
        file = open(self.filename, 'w')
        try:
            yaml.dump(self.root.raw(), file)
        finally:
            file.close()

    def OnClose(self, event):
        self.Close()

    def OnTreeChange(self, event):
        self.graph.Refresh()

    def OnTreeSelected(self, event):
        self.graph.selected = self._tree.selected

    def OnGraphSelected(self, event):
        self._tree.selected = event.target

class Graph(wx.ScrolledWindow):
    def __init__(self, root, radius, padding, *args, **kwargs):
        wx.ScrolledWindow.__init__(self, *args, **kwargs)

        self.radius = radius
        self.padding = padding
        self._selected = None

        self.root = root
        self.Refresh()
        self.Bind(wx.EVT_PAINT, self.Redraw)
        self.Bind(wx.EVT_LEFT_DOWN, self.OnMouseClick)

    def getselected(self):
        return self._selected
    def setselected(self, value):
        if self._selected:
            # Restore previous selected
            self.DrawNode(self._selected)

        self._selected = value

        if self._selected:
            self.DrawNode(self._selected, bgcolor='#FFFF80')
    selected = property(getselected, setselected)

    def Refresh(self):
        self.selected = None
        graph = pyggdrasil.graph.generate(self.root)

        scalar = 2 * (self.radius + self.padding)
        self.nodes = dict((node, pos * scalar)
                                     for (node, pos) in graph)

        pos = self.GetViewStart()
        self.SetScrollbars(1, 1, graph.width * scalar, graph.height * scalar)
        self.Scroll(*pos)

    def Redraw(self, event=None):
        dc = wx.PaintDC(self)
        self.DoPrepareDC(dc)
        dc.Clear()

        lines = []
        for node in self.nodes:
            if node.parent:
                npos = self.nodes[node]
                ppos = self.nodes[node.parent]
                vector = ppos - npos
                magnitude = abs(vector)
                direction = math.atan2(vector.imag, vector.real)

                # Adjust for the circles
                ppos -= self.radius * cmath.exp(direction * 1j)
                npos += self.radius * cmath.exp(direction * 1j)

                # Relational line
                lines.append((npos.real, npos.imag, ppos.real, ppos.imag))

                # Little arrow at the end
                pos1 = ppos - self.padding * cmath.exp((direction + 0.5) * 1j)
                pos2 = ppos - self.padding * cmath.exp((direction - 0.5) * 1j)
                lines.append((pos1.real, pos1.imag, ppos.real, ppos.imag))
                lines.append((pos2.real, pos2.imag, ppos.real, ppos.imag))

        if lines:
            dc.DrawLineList(lines)

        for node in self.nodes:
            self.DrawNode(node, dc)

        if self.selected:
            self.DrawNode(self.selected, dc=dc, bgcolor='#FFFF80')

        dc.EndDrawing()

    def DrawNode(self, node, dc=None, bgcolor=None):
        if dc == None:
            dc = wx.ClientDC(self)
            self.DoPrepareDC(dc)
            enddrawing = True
        else:
            enddrawing = False

        if bgcolor:
            dc.SetBrush(wx.Brush(bgcolor))

        pos = self.nodes[node]
        dc.DrawCircle(pos.real, pos.imag, self.radius)

        w, h = dc.GetTextExtent(node.id)
        dc.DrawText(node.id, pos.real - w/2.0, pos.imag - h/2.0)

        if enddrawing:
            dc.EndDrawing()

    def OnMouseClick(self, event):
        dc = wx.ClientDC(self)
        self.DoPrepareDC(dc)
        click = tuple(event.GetLogicalPosition(dc))

        for (node, pos) in self.nodes.items():
            if (pos.real - self.radius) <= click[0] <= (pos.real + self.radius) and \
               (pos.imag - self.radius) <= click[1] <= (pos.imag + self.radius):
                wx.PostEvent(self, GraphSelectedEvent(target=node))
                return

class Tree(wx.Panel):
    def __init__(self, root, *args, **kwargs):
        wx.Panel.__init__(self, *args, **kwargs)

        sizer = wx.BoxSizer(wx.VERTICAL)

        self._tree = wx.TreeCtrl(self, wx.ID_ANY,
                                style=(wx.TR_EDIT_LABELS | wx.TR_HAS_BUTTONS))
        self._tree.Bind(wx.EVT_TREE_END_LABEL_EDIT, self.OnRename)
        self._tree.Bind(wx.EVT_TREE_BEGIN_DRAG, self.OnBeginDrag)
        self._tree.Bind(wx.EVT_TREE_END_DRAG, self.OnEndDrag)
        sizer.Add(self._tree, 1, wx.EXPAND)

        sizer.Add(self._buttoncontrols())

        self.SetSizer(sizer)

        self.root = root
        self.Refresh()

    def _buttoncontrols(self):
        sizer = wx.FlexGridSizer(rows=0, cols=2)

        self._childinput = wx.TextCtrl(self, wx.ID_ANY, style=wx.TE_PROCESS_ENTER)
        self._childinput.Bind(wx.EVT_TEXT_ENTER, self.OnAdd)
        sizer.Add(self._childinput, 1, wx.EXPAND)

        add = wx.Button(self, wx.ID_ADD, 'Add Child')
        add.Bind(wx.EVT_BUTTON, self.OnAdd)
        sizer.Add(add, 1, wx.EXPAND)

        sizer.AddStretchSpacer()

        remove = wx.Button(self, wx.ID_REMOVE)
        remove.Bind(wx.EVT_BUTTON, self.OnRemove)
        sizer.Add(remove, 1, wx.EXPAND)

        sort = wx.Button(self, wx.ID_ANY, 'Sort')
        sort.Bind(wx.EVT_BUTTON, self.OnSort)
        sizer.Add(sort, 1, wx.EXPAND)

        self._autosort = wx.CheckBox(self, wx.ID_ANY, 'Auto Sort')
        self._autosort.Bind(wx.EVT_CHECKBOX, self.OnSort)
        sizer.Add(self._autosort)

        return sizer

    def getselected(self):
        return self.nodes[self._tree.GetSelection()]
    def setselected(self, value):
        item = self.nodes.getkey(value)
        self._tree.SelectItem(item)
    selected = property(getselected, setselected)

    @property
    def autosort(self):
        return self._autosort.IsChecked()

    def Refresh(self):
        self._tree.DeleteAllItems()
        self.nodes = pyggdrasil.model.EqualsDict()
        self._PopulateTreeItem(self.root, None)

    def _PopulateTreeItem(self, node, parent):
        if parent:
            treeitem = self._tree.AppendItem(parent, node.id)
        else:
            treeitem = self._tree.AddRoot(node.id)
        self.nodes[treeitem] = node

        for child in node.children:
            self._PopulateTreeItem(child, treeitem)

    def _SortItem(self, item):
        self._tree.SortChildren(item)
        self.nodes[item].sort()

        # I want an iterator
        child = self._tree.GetFirstChild(item)[0]
        while child.IsOk():
            self._SortItem(child)
            child = self._tree.GetNextSibling(child)

    def OnSort(self, event):
        self._SortItem(self._tree.GetRootItem())
        wx.PostEvent(self, TreeChangedEvent())

    def OnRename(self, event):
        selectedid = self._tree.GetSelection()
        nodeid = event.GetLabel()
        self.nodes[selectedid].id = str(nodeid)

        wx.PostEvent(self, TreeChangedEvent())

    def OnAdd(self, event):
        selectedid = self._tree.GetSelection()
        nodeid = self._childinput.GetValue()

        if nodeid:
            node = pyggdrasil.model.Node(str(nodeid), None, self.nodes[selectedid])

            newid = self._tree.AppendItem(selectedid, nodeid)
            self.nodes[newid] = node
            self._tree.Expand(selectedid)

            self._childinput.Clear()

            if self.autosort:
                self._SortItem(selectedid)
            wx.PostEvent(self, TreeChangedEvent())

    def OnRemove(self, event):
        #TODO: Deal with children somehow
        selectedid = self._tree.GetSelection()

        # Do not let root node get removed
        if self.nodes[selectedid].parent:
            self.nodes[selectedid].parent = None
            del self.nodes[selectedid]

            self._tree.Delete(selectedid)

            wx.PostEvent(self, TreeChangedEvent())

    def OnBeginDrag(self, event):
        self._dragitem = event.GetItem()
        if self._dragitem.IsOk():
            event.Allow()

    def OnEndDrag(self, event):
        parent = event.GetItem()

        if not parent.IsOk():
            return
        if parent == self._tree.GetItemParent(self._dragitem):
            return
        try:
            oldid = self._dragitem
        except AttributeError:
            return

        text = self._tree.GetItemText(oldid)
        self._tree.Delete(oldid)

        node = self.nodes.pop(oldid)
        node.parent = self.nodes[parent]
        self._PopulateTreeItem(node, parent)
        self._tree.Expand(parent)

        wx.PostEvent(self, TreeChangedEvent())

class App(wx.App):
    def OnInit(self):
        self.SetAppName('Pyggdrasil')
        root = pyggdrasil.model.Node('root', None)
        frame = Main(parent=None, id=wx.ID_ANY)
        frame.Bind(wx.EVT_MENU, self.OnMenuExit, id=wx.ID_EXIT)
        frame.Show(True)
        return True

    def OnMenuExit(self, event):
        self.ExitMainLoop()
<commit_msg>Renamed variables to be more consistent<commit_after>import wx
import wx.lib.newevent
import math
import cmath
import yaml

import pyggdrasil

TreeChangedEvent, TREE_CHANGED_EVENT = wx.lib.newevent.NewEvent()
GraphSelectedEvent, GRAPH_SELECTED_EVENT = wx.lib.newevent.NewEvent()

class Main(wx.Frame):
    def __init__(self, root=None, filename=None, *args, **kwargs):
        wx.Frame.__init__(self, *args, **kwargs)

        self.root = root or pyggdrasil.model.Node('root', None)
        self.filename = filename

        menubar = wx.MenuBar()

        file = wx.Menu()
        new = file.Append(wx.ID_NEW)
        self.Bind(wx.EVT_MENU, self.OnNew, new)
        open = file.Append(wx.ID_OPEN)
        self.Bind(wx.EVT_MENU, self.OnOpen, open)
        file.AppendSeparator()

        close = file.Append(wx.ID_CLOSE, 'Close\tCtrl-W')
        self.Bind(wx.EVT_MENU, self.OnClose, close)
        save = file.Append(wx.ID_SAVE)
        self.Bind(wx.EVT_MENU, self.OnSave, save)
        saveas = file.Append(wx.ID_SAVEAS)
        self.Bind(wx.EVT_MENU, self.OnSaveAs, saveas)
        file.AppendSeparator()

        file.Append(wx.ID_EXIT)
        menubar.Append(file, '&File')

        self.SetMenuBar(menubar)

        box = wx.BoxSizer(wx.HORIZONTAL)

        self._tree = Tree(self.root, self)
        box.Add(self._tree, 0, wx.EXPAND)

        self.graph = Graph(self.root, 40, 5, self)
        box.Add(self.graph, 1, wx.EXPAND)

        self.SetSizer(box)

        self._tree.Bind(TREE_CHANGED_EVENT, self.OnTreeChange)
        self._tree.Bind(wx.EVT_TREE_SEL_CHANGED, self.OnTreeSelected)
        self.graph.Bind(GRAPH_SELECTED_EVENT, self.OnGraphSelected)

    def getfilename(self):
        return self._filename
    def setfilename(self, value):
        self._filename = value
        self.SetTitle('Pyggdrasil - ' + (self.filename or 'new'))
    filename = property(getfilename, setfilename)

    def OnNew(self, event):
        frame = Main(parent=self.GetParent(), id=wx.ID_ANY)
        frame.Show(True)

    def OnOpen(self, event):
        filename = wx.LoadFileSelector('Pyggdrasil', extension='pyg', parent=self)
        if filename:
            file = open(filename)
            try:
                root = pyggdrasil.model.Node.from_raw(yaml.load(file))
                frame = Main(root, filename, self.GetParent(), wx.ID_ANY)
                frame.Show(True)
            finally:
                file.close()

    def OnSave(self, event):
        if self.filename:
            self._Save()
        else:
            self.OnSaveAs(event)

    def OnSaveAs(self, event):
        filename = wx.SaveFileSelector('Pyggdrasil', extension='pyg', parent=self)
        if filename:
            if '.' not in filename:
                filename += '.pyg'
            self.filename = filename
            self._Save()

    def _Save(self):
        file = open(self.filename, 'w')
        try:
            yaml.dump(self.root.raw(), file)
        finally:
            file.close()

    def OnClose(self, event):
        self.Close()

    def OnTreeChange(self, event):
        self.graph.Refresh()

    def OnTreeSelected(self, event):
        self.graph.selected = self._tree.selected

    def OnGraphSelected(self, event):
        self._tree.selected = event.target

class Graph(wx.ScrolledWindow):
    def __init__(self, root, radius, padding, *args, **kwargs):
        wx.ScrolledWindow.__init__(self, *args, **kwargs)

        self.radius = radius
        self.padding = padding
        self._selected = None

        self.root = root
        self.Refresh()
        self.Bind(wx.EVT_PAINT, self.Redraw)
        self.Bind(wx.EVT_LEFT_DOWN, self.OnMouseClick)

    def getselected(self):
        return self._selected
    def setselected(self, value):
        if self._selected:
            # Restore previous selected
            self.DrawNode(self._selected)

        self._selected = value

        if self._selected:
            self.DrawNode(self._selected, bgcolor='#FFFF80')
    selected = property(getselected, setselected)

    def Refresh(self):
        self.selected = None
        graph = pyggdrasil.graph.generate(self.root)

        scalar = 2 * (self.radius + self.padding)
        self.nodes = dict((node, pos * scalar)
                                     for (node, pos) in graph)

        pos = self.GetViewStart()
        self.SetScrollbars(1, 1, graph.width * scalar, graph.height * scalar)
        self.Scroll(*pos)

    def Redraw(self, event=None):
        dc = wx.PaintDC(self)
        self.DoPrepareDC(dc)
        dc.Clear()

        lines = []
        for node in self.nodes:
            if node.parent:
                npos = self.nodes[node]
                ppos = self.nodes[node.parent]
                vector = ppos - npos
                magnitude = abs(vector)
                direction = math.atan2(vector.imag, vector.real)

                # Adjust for the circles
                ppos -= self.radius * cmath.exp(direction * 1j)
                npos += self.radius * cmath.exp(direction * 1j)

                # Relational line
                lines.append((npos.real, npos.imag, ppos.real, ppos.imag))

                # Little arrow at the end
                pos1 = ppos - self.padding * cmath.exp((direction + 0.5) * 1j)
                pos2 = ppos - self.padding * cmath.exp((direction - 0.5) * 1j)
                lines.append((pos1.real, pos1.imag, ppos.real, ppos.imag))
                lines.append((pos2.real, pos2.imag, ppos.real, ppos.imag))

        if lines:
            dc.DrawLineList(lines)

        for node in self.nodes:
            self.DrawNode(node, dc)

        if self.selected:
            self.DrawNode(self.selected, dc=dc, bgcolor='#FFFF80')

        dc.EndDrawing()

    def DrawNode(self, node, dc=None, bgcolor=None):
        if dc == None:
            dc = wx.ClientDC(self)
            self.DoPrepareDC(dc)
            enddrawing = True
        else:
            enddrawing = False

        if bgcolor:
            dc.SetBrush(wx.Brush(bgcolor))

        pos = self.nodes[node]
        dc.DrawCircle(pos.real, pos.imag, self.radius)

        w, h = dc.GetTextExtent(node.id)
        dc.DrawText(node.id, pos.real - w/2.0, pos.imag - h/2.0)

        if enddrawing:
            dc.EndDrawing()

    def OnMouseClick(self, event):
        dc = wx.ClientDC(self)
        self.DoPrepareDC(dc)
        click = tuple(event.GetLogicalPosition(dc))

        for (node, pos) in self.nodes.items():
            if (pos.real - self.radius) <= click[0] <= (pos.real + self.radius) and \
               (pos.imag - self.radius) <= click[1] <= (pos.imag + self.radius):
                wx.PostEvent(self, GraphSelectedEvent(target=node))
                return

class Tree(wx.Panel):
    def __init__(self, root, *args, **kwargs):
        wx.Panel.__init__(self, *args, **kwargs)

        sizer = wx.BoxSizer(wx.VERTICAL)

        self._tree = wx.TreeCtrl(self, wx.ID_ANY,
                                style=(wx.TR_EDIT_LABELS | wx.TR_HAS_BUTTONS))
        self._tree.Bind(wx.EVT_TREE_END_LABEL_EDIT, self.OnRename)
        self._tree.Bind(wx.EVT_TREE_BEGIN_DRAG, self.OnBeginDrag)
        self._tree.Bind(wx.EVT_TREE_END_DRAG, self.OnEndDrag)
        sizer.Add(self._tree, 1, wx.EXPAND)

        sizer.Add(self._buttoncontrols())

        self.SetSizer(sizer)

        self.root = root
        self.Refresh()

    def _buttoncontrols(self):
        sizer = wx.FlexGridSizer(rows=0, cols=2)

        self._childinput = wx.TextCtrl(self, wx.ID_ANY, style=wx.TE_PROCESS_ENTER)
        self._childinput.Bind(wx.EVT_TEXT_ENTER, self.OnAdd)
        sizer.Add(self._childinput, 1, wx.EXPAND)

        add = wx.Button(self, wx.ID_ADD, 'Add Child')
        add.Bind(wx.EVT_BUTTON, self.OnAdd)
        sizer.Add(add, 1, wx.EXPAND)

        sizer.AddStretchSpacer()

        remove = wx.Button(self, wx.ID_REMOVE)
        remove.Bind(wx.EVT_BUTTON, self.OnRemove)
        sizer.Add(remove, 1, wx.EXPAND)

        sort = wx.Button(self, wx.ID_ANY, 'Sort')
        sort.Bind(wx.EVT_BUTTON, self.OnSort)
        sizer.Add(sort, 1, wx.EXPAND)

        self._autosort = wx.CheckBox(self, wx.ID_ANY, 'Auto Sort')
        self._autosort.Bind(wx.EVT_CHECKBOX, self.OnSort)
        sizer.Add(self._autosort)

        return sizer

    def getselected(self):
        return self.nodes[self._tree.GetSelection()]
    def setselected(self, value):
        item = self.nodes.getkey(value)
        self._tree.SelectItem(item)
    selected = property(getselected, setselected)

    @property
    def autosort(self):
        return self._autosort.IsChecked()

    def Refresh(self):
        self._tree.DeleteAllItems()
        self.nodes = pyggdrasil.model.EqualsDict()
        self._PopulateTreeItem(self.root, None)

    def _PopulateTreeItem(self, node, parent):
        if parent:
            treeitem = self._tree.AppendItem(parent, node.id)
        else:
            treeitem = self._tree.AddRoot(node.id)
        self.nodes[treeitem] = node

        for child in node.children:
            self._PopulateTreeItem(child, treeitem)

    def _SortItem(self, item):
        self._tree.SortChildren(item)
        self.nodes[item].sort()

        # I want an iterator
        child = self._tree.GetFirstChild(item)[0]
        while child.IsOk():
            self._SortItem(child)
            child = self._tree.GetNextSibling(child)

    def OnSort(self, event):
        self._SortItem(self._tree.GetRootItem())
        wx.PostEvent(self, TreeChangedEvent())

    def OnRename(self, event):
        item = self._tree.GetSelection()
        nodeid = event.GetLabel()
        self.nodes[item].id = str(nodeid)

        wx.PostEvent(self, TreeChangedEvent())

    def OnAdd(self, event):
        item = self._tree.GetSelection()
        nodeid = self._childinput.GetValue()

        if nodeid:
            node = pyggdrasil.model.Node(str(nodeid), None, self.nodes[item])

            newitem = self._tree.AppendItem(item, nodeid)
            self.nodes[newitem] = node
            self._tree.Expand(item)

            self._childinput.Clear()

            if self.autosort:
                self._SortItem(item)
            wx.PostEvent(self, TreeChangedEvent())

    def OnRemove(self, event):
        #TODO: Deal with children somehow
        item = self._tree.GetSelection()

        # Do not let root node get removed
        if self.nodes[item].parent:
            self.nodes[item].parent = None
            del self.nodes[item]

            self._tree.Delete(item)

            wx.PostEvent(self, TreeChangedEvent())

    def OnBeginDrag(self, event):
        self._dragitem = event.GetItem()
        if self._dragitem.IsOk():
            event.Allow()

    def OnEndDrag(self, event):
        parent = event.GetItem()

        if not parent.IsOk():
            return
        if parent == self._tree.GetItemParent(self._dragitem):
            return
        try:
            oldid = self._dragitem
        except AttributeError:
            return

        text = self._tree.GetItemText(oldid)
        self._tree.Delete(oldid)

        node = self.nodes.pop(oldid)
        node.parent = self.nodes[parent]
        self._PopulateTreeItem(node, parent)
        self._tree.Expand(parent)

        wx.PostEvent(self, TreeChangedEvent())

class App(wx.App):
    def OnInit(self):
        self.SetAppName('Pyggdrasil')
        root = pyggdrasil.model.Node('root', None)
        frame = Main(parent=None, id=wx.ID_ANY)
        frame.Bind(wx.EVT_MENU, self.OnMenuExit, id=wx.ID_EXIT)
        frame.Show(True)
        return True

    def OnMenuExit(self, event):
        self.ExitMainLoop()
<|endoftext|>


<commit_before>import argparse
import sys

from . import pygs
from . import config
from . import core

# An extension of argparse as defined by https://stackoverflow.com/a/26379693/741316
def set_default_subparser(self, name, args=None):
    """default subparser selection. Call after setup, just before parse_args()
    name: is the name of the subparser to call by default
    args: if set is the argument list handed to parse_args()

    , tested with 2.7, 3.2, 3.3, 3.4
    it works with 2.6 assuming argparse is installed
    """
    subparser_found = False
    for arg in sys.argv[1:]:
        if arg in ['-h', '--help']:  # global help if no subparser
            break
    else:
        for x in self._subparsers._actions:
            if not isinstance(x, argparse._SubParsersAction):
                continue
            for sp_name in x._name_parser_map.keys():
                if sp_name in sys.argv[1:]:
                    subparser_found = True
        if not subparser_found:
            # insert default in first position, this implies no
            # global options without a sub_parsers specified
            if args is None:
                sys.argv.insert(1, name)
            else:
                args.insert(0, name)

argparse.ArgumentParser.set_default_subparser = set_default_subparser

class RepeatedOptionWithDefault(argparse.Action):
    # https://bugs.python.org/issue16399
    def __call__(self, parser, namespace, value, option_string=None):
        dest = getattr(namespace, self.dest, [])
        if dest is self.default:
            dest = []
        dest.append(value)
        setattr(namespace, self.dest, dest)

def accounts_cat(infile, accounts_file):
    from . import wizard as wizard_mod

    if infile is None or infile.isatty():
        # pyggybank accounts > foo.txt
        # pyggybank accounts

        infile = infile or sys.stdout
        sys.exit(wizard_mod.cat_account(infile, accounts_file=accounts_file))
    else:
        # pyggybank accounts < foo.txt
        content = infile.read()
        sys.exit(wizard_mod.write_account(content, accounts_file=accounts_file))

def main():
    parser = argparse.ArgumentParser(description='Get information from your internet banks')

    top_subparser = parser.add_subparsers()
    sp = top_subparser.add_parser('balance')
    sp.set_defaults(func=balance)
    wizard_parser = top_subparser.add_parser('wizard')
    wizard_parser.set_defaults(func=wizard)
    accounts_parser = top_subparser.add_parser('accounts')
    import sys
    accounts_parser.add_argument('infile', nargs='?', type=argparse.FileType('r'),
                                 default=sys.stdin if not sys.stdin.isatty() else None)

    accounts_parser.set_defaults(func=accounts_cat)

    # TODO: Should this just be controled by config, and not a CLI arg...?
    wizard_parser.add_argument('--gpg-home', dest='gpg_home', default=None)
    accounts_parser.add_argument('--gpg-home', dest='gpg_home', default=None)

    accounts_parser.add_argument(
            '--accounts-file',
            default=pygs.DEFAULT_ACCOUNTS_FILE,
            dest='accounts_file',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    wizard_parser.add_argument(
            '--accounts-file',
            default=pygs.DEFAULT_ACCOUNTS_FILE,
            dest='accounts_file',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    sp.add_argument(
            '--accounts',
            action=RepeatedOptionWithDefault,
            default=[pygs.DEFAULT_ACCOUNTS_FILE],
            dest='accounts_files',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    # Use the extension defined above.
    parser.set_default_subparser('balance')
    parsed_args = vars(parser.parse_args())

    gpg_home = parsed_args.pop('gpg_home', None)
    config.CONFIG['gpg-home'] = gpg_home

    fn = parsed_args.pop('func')
    sys.exit(fn(**parsed_args))

def balance(accounts_files):
    # Q: Do we really need to support multiple configs at the same time?
    from . import gpgconfig
    from pathlib import Path
    import splinter
    browser = splinter.Browser(config.CONFIG.get('browser', 'chrome'))
    try:
        for accounts_file in accounts_files:
            accounts_file = Path(accounts_file)
            if not accounts_file.exists():
                raise ValueError("The accounts file doesn't exist. Perhaps you "
                                 "want to run the wizard?")
            accounts_config = gpgconfig.decrypt_config(accounts_file)
            for account in accounts_config.get('accounts', []):
                provider, credentials = core.Provider.from_config(account)
                provider.authenticate(browser, credentials)
                # for account_balance in provider.balances():
                #     print(account_balance.name, account_balance.total)
    finally:
        browser.quit()

def wizard(*args, **kwargs):
    from . import wizard as wizard_mod
    sys.exit(wizard_mod.wizard(*args, **kwargs))

if __name__ == '__main__':
    main()
<commit_msg>Added the 'pyggybank visit' feature to allow users to open their internet banking pages without having to go through the authentication manually.<commit_after>import argparse
import sys

import prompt_toolkit

from . import pygs
from . import config
from . import core

# An extension of argparse as defined by https://stackoverflow.com/a/26379693/741316
def set_default_subparser(self, name, args=None):
    """default subparser selection. Call after setup, just before parse_args()
    name: is the name of the subparser to call by default
    args: if set is the argument list handed to parse_args()

    , tested with 2.7, 3.2, 3.3, 3.4
    it works with 2.6 assuming argparse is installed
    """
    subparser_found = False
    for arg in sys.argv[1:]:
        if arg in ['-h', '--help']:  # global help if no subparser
            break
    else:
        for x in self._subparsers._actions:
            if not isinstance(x, argparse._SubParsersAction):
                continue
            for sp_name in x._name_parser_map.keys():
                if sp_name in sys.argv[1:]:
                    subparser_found = True
        if not subparser_found:
            # insert default in first position, this implies no
            # global options without a sub_parsers specified
            if args is None:
                sys.argv.insert(1, name)
            else:
                args.insert(0, name)

argparse.ArgumentParser.set_default_subparser = set_default_subparser

class RepeatedOptionWithDefault(argparse.Action):
    # https://bugs.python.org/issue16399
    def __call__(self, parser, namespace, value, option_string=None):
        dest = getattr(namespace, self.dest, [])
        if dest is self.default:
            dest = []
        dest.append(value)
        setattr(namespace, self.dest, dest)

def accounts_cat(infile, accounts_file):
    from . import wizard as wizard_mod

    if infile is None or infile.isatty():
        # pyggybank accounts > foo.txt
        # pyggybank accounts

        infile = infile or sys.stdout
        sys.exit(wizard_mod.cat_account(infile, accounts_file=accounts_file))
    else:
        # pyggybank accounts < foo.txt
        content = infile.read()
        sys.exit(wizard_mod.write_account(content, accounts_file=accounts_file))

def main():
    parser = argparse.ArgumentParser(description='Get information from your internet banks')

    top_subparser = parser.add_subparsers()
    sp = top_subparser.add_parser('balance')
    sp.set_defaults(func=balance)

    visit_parser = top_subparser.add_parser('visit')
    visit_parser.set_defaults(func=visit)

    wizard_parser = top_subparser.add_parser('wizard')
    wizard_parser.set_defaults(func=wizard)
    accounts_parser = top_subparser.add_parser('accounts')
    import sys
    accounts_parser.add_argument('infile', nargs='?', type=argparse.FileType('r'),
                                 default=sys.stdin if not sys.stdin.isatty() else None)

    accounts_parser.set_defaults(func=accounts_cat)

    # TODO: Should this just be controled by config, and not a CLI arg...?
    wizard_parser.add_argument('--gpg-home', dest='gpg_home', default=None)
    accounts_parser.add_argument('--gpg-home', dest='gpg_home', default=None)

    accounts_parser.add_argument(
            '--accounts-file',
            default=pygs.DEFAULT_ACCOUNTS_FILE,
            dest='accounts_file',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    wizard_parser.add_argument(
            '--accounts-file',
            default=pygs.DEFAULT_ACCOUNTS_FILE,
            dest='accounts_file',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    sp.add_argument(
            '--accounts',
            action=RepeatedOptionWithDefault,
            default=[pygs.DEFAULT_ACCOUNTS_FILE],
            dest='accounts_files',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    visit_parser.add_argument(
            '--accounts',
            action=RepeatedOptionWithDefault,
            default=[pygs.DEFAULT_ACCOUNTS_FILE],
            dest='accounts_files',
            help=('The accounts file to use (default: {})'
                  ''.format(pygs.DEFAULT_ACCOUNTS_FILE)))

    # Use the extension defined above.
    parser.set_default_subparser('balance')
    parsed_args = vars(parser.parse_args())

    gpg_home = parsed_args.pop('gpg_home', None)
    config.CONFIG['gpg-home'] = gpg_home

    fn = parsed_args.pop('func')
    sys.exit(fn(**parsed_args))

def balance(accounts_files):
    # Q: Do we really need to support multiple configs at the same time?
    from . import gpgconfig
    from pathlib import Path
    import splinter
    browser = splinter.Browser(config.CONFIG.get('browser', 'chrome'))
    try:
        for accounts_file in accounts_files:
            accounts_file = Path(accounts_file)
            if not accounts_file.exists():
                raise ValueError("The accounts file doesn't exist. Perhaps you "
                                 "want to run the wizard?")
            accounts_config = gpgconfig.decrypt_config(accounts_file)
            for account in accounts_config.get('accounts', []):
                provider, credentials = core.Provider.from_config(account)
                provider.authenticate(browser, credentials)

                # for account_balance in provider.balances():
                #     print(account_balance.name, account_balance.total)
    finally:
        browser.quit()

def wizard(*args, **kwargs):
    from . import wizard as wizard_mod
    sys.exit(wizard_mod.wizard(*args, **kwargs))

def visit(accounts_files):
    # Q: Do we really need to support multiple configs at the same time?
    from . import gpgconfig
    from pathlib import Path
    import splinter
    browser = splinter.Browser(config.CONFIG.get('browser', 'chrome'))

    try:
        account_tabs = []
        for accounts_file in accounts_files:
            accounts_file = Path(accounts_file)
            if not accounts_file.exists():
                raise ValueError("The accounts file doesn't exist. Perhaps you "
                                 "want to run the wizard?")
            accounts_config = gpgconfig.decrypt_config(accounts_file)
            for account in accounts_config.get('accounts', []):
                provider, credentials = core.Provider.from_config(account)
                browser.execute_script('''window.open("about:blank", "_blank");''')
                window = browser.windows[-1]
                window.is_current = True
                steps = provider.authenticate(browser, credentials)
                account_tabs.append([window, steps])
                # for account_balance in provider.balances():
                #     print(account_balance.name, account_balance.total)
        while account_tabs:
            for window_and_steps in account_tabs[:]:
                window, steps = window_and_steps
                window.is_current = True
                try:
                    next(steps)
                except StopIteration:
                    account_tabs.remove(window_and_steps)

        while prompt_toolkit.prompt('Ready to close the browser? ').lower() in ['no', 'n']:
            pass
    finally:
        browser.quit()

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


<commit_before># -*- coding: utf-8 -*-
"""
    pygments.util
    ~~~~~~~~~~~~~

    Utility functions.

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import re
import sys

split_path_re = re.compile(r'[/\\ ]')
doctype_lookup_re = re.compile(r'''(?smx)
    (<\?.*?\?>)?\s*
    <!DOCTYPE\s+(
     [a-zA-Z_][a-zA-Z0-9]*\s+
     [a-zA-Z_][a-zA-Z0-9]*\s+
     "[^"]*")
     [^>]*>
''')
tag_re = re.compile(r'<(.+?)(\s.*?)?>.*?</.+?>(?uism)')

class ClassNotFound(ValueError):
    """
    If one of the get_*_by_* functions didn't find a matching class.
    """

class OptionError(Exception):
    pass

def get_choice_opt(options, optname, allowed, default=None, normcase=False):
    string = options.get(optname, default)
    if normcase:
        string = string.lower()
    if string not in allowed:
        raise OptionError('Value for option %s must be one of %s' %
                          (optname, ', '.join(map(str, allowed))))
    return string

def get_bool_opt(options, optname, default=None):
    string = options.get(optname, default)
    if isinstance(string, bool):
        return string
    elif isinstance(string, int):
        return bool(string)
    elif not isinstance(string, string_types):
        raise OptionError('Invalid type %r for option %s; use '
                          '1/0, yes/no, true/false, on/off' % (
                          string, optname))
    elif string.lower() in ('1', 'yes', 'true', 'on'):
        return True
    elif string.lower() in ('0', 'no', 'false', 'off'):
        return False
    else:
        raise OptionError('Invalid value %r for option %s; use '
                          '1/0, yes/no, true/false, on/off' % (
                          string, optname))

def get_int_opt(options, optname, default=None):
    string = options.get(optname, default)
    try:
        return int(string)
    except TypeError:
        raise OptionError('Invalid type %r for option %s; you '
                          'must give an integer value' % (
                          string, optname))
    except ValueError:
        raise OptionError('Invalid value %r for option %s; you '
                          'must give an integer value' % (
                          string, optname))

def get_list_opt(options, optname, default=None):
    val = options.get(optname, default)
    if isinstance(val, string_types):
        return val.split()
    elif isinstance(val, (list, tuple)):
        return list(val)
    else:
        raise OptionError('Invalid type %r for option %s; you '
                          'must give a list value' % (
                          val, optname))

def docstring_headline(obj):
    if not obj.__doc__:
        return ''
    res = []
    for line in obj.__doc__.strip().splitlines():
        if line.strip():
            res.append(" " + line.strip())
        else:
            break
    return ''.join(res).lstrip()

def make_analysator(f):
    """
    Return a static text analysation function that
    returns float values.
    """
    def text_analyse(text):
        try:
            rv = f(text)
        except Exception:
            return 0.0
        if not rv:
            return 0.0
        try:
            return min(1.0, max(0.0, float(rv)))
        except (ValueError, TypeError):
            return 0.0
    text_analyse.__doc__ = f.__doc__
    return staticmethod(text_analyse)

def shebang_matches(text, regex):
    """
    Check if the given regular expression matches the last part of the
    shebang if one exists.

        >>> from pygments.util import shebang_matches
        >>> shebang_matches('#!/usr/bin/env python', r'python(2\.\d)?')
        True
        >>> shebang_matches('#!/usr/bin/python2.4', r'python(2\.\d)?')
        True
        >>> shebang_matches('#!/usr/bin/python-ruby', r'python(2\.\d)?')
        False
        >>> shebang_matches('#!/usr/bin/python/ruby', r'python(2\.\d)?')
        False
        >>> shebang_matches('#!/usr/bin/startsomethingwith python',
        ...                 r'python(2\.\d)?')
        True

    It also checks for common windows executable file extensions::

        >>> shebang_matches('#!C:\\Python2.4\\Python.exe', r'python(2\.\d)?')
        True

    Parameters (``'-f'`` or ``'--foo'`` are ignored so ``'perl'`` does
    the same as ``'perl -e'``)

    Note that this method automatically searches the whole string (eg:
    the regular expression is wrapped in ``'^$'``)
    """
    index = text.find('\n')
    if index >= 0:
        first_line = text[:index].lower()
    else:
        first_line = text.lower()
    if first_line.startswith('#!'):
        try:
            found = [x for x in split_path_re.split(first_line[2:].strip())
                     if x and not x.startswith('-')][-1]
        except IndexError:
            return False
        regex = re.compile('^%s(\.(exe|cmd|bat|bin))?$' % regex, re.IGNORECASE)
        if regex.search(found) is not None:
            return True
    return False

def doctype_matches(text, regex):
    """
    Check if the doctype matches a regular expression (if present).
    Note that this method only checks the first part of a DOCTYPE.
    eg: 'html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN"'
    """
    m = doctype_lookup_re.match(text)
    if m is None:
        return False
    doctype = m.group(2)
    return re.compile(regex).match(doctype.strip()) is not None

def html_doctype_matches(text):
    """
    Check if the file looks like it has a html doctype.
    """
    return doctype_matches(text, r'html\s+PUBLIC\s+"-//W3C//DTD X?HTML.*')

_looks_like_xml_cache = {}
def looks_like_xml(text):
    """
    Check if a doctype exists or if we have some tags.
    """
    key = hash(text)
    try:
        return _looks_like_xml_cache[key]
    except KeyError:
        m = doctype_lookup_re.match(text)
        if m is not None:
            return True
        rv = tag_re.search(text[:1000]) is not None
        _looks_like_xml_cache[key] = rv
        return rv

# Python narrow build compatibility

def _surrogatepair(c):
    # Given a unicode character code
    # with length greater than 16 bits,
    # print the two 16 bit surrogate pair.
    # From example D28 of:
    # http://www.unicode.org/book/ch03.pdf
    return (0xd7c0 + (c >> 10), (0xdc00 + (c & 0x3ff)))

def unirange(a, b):
    """
    Returns a regular expression string to match the given non-BMP range.
    """
    if b < a:
        raise ValueError("Bad character range")
    if a < 0x10000 or b < 0x10000:
        raise ValueError("unirange is only defined for non-BMP ranges")

    if sys.maxunicode > 0xffff:
        # wide build
        return u'[%s-%s]' % (unichr(a), unichr(b))
    else:
        # narrow build stores surrogates, and the 're' module handles them
        # (incorrectly) as characters.  Since there is still ordering among
        # these characters, expand the range to one that it understands.  Some
        # background in http://bugs.python.org/issue3665 and
        # http://bugs.python.org/issue12749
        #
        # Additionally, the lower constants are using unichr rather than
        # literals because jython [which uses the wide path] can't load this
        # file if they are literals.
        ah, al = _surrogatepair(a)
        bh, bl = _surrogatepair(b)
        if ah == bh:
            return u'(?:%s[%s-%s])' % (unichr(ah), unichr(al), unichr(bl))
        else:
            buf = []
            buf.append(u'%s[%s-%s]' %
                       (unichr(ah), unichr(al),
                        ah == bh and unichr(bl) or unichr(0xdfff)))
            if ah - bh > 1:
                buf.append(u'[%s-%s][%s-%s]' %
                           unichr(ah+1), unichr(bh-1), unichr(0xdc00), unichr(0xdfff))
            if ah != bh:
                buf.append(u'%s[%s-%s]' %
                           (unichr(bh), unichr(0xdc00), unichr(bl)))

            return u'(?:' + u'|'.join(buf) + u')'

# Python 2/3 compatibility

if sys.version_info < (3, 0):
    unichr = unichr
    xrange = xrange
    string_types = (str, unicode)
    text_type = unicode
    u_prefix = 'u'
    iteritems = dict.iteritems
    itervalues = dict.itervalues
    import StringIO, cStringIO
    # unfortunately, io.StringIO in Python 2 doesn't accept str at all
    StringIO = StringIO.StringIO
    BytesIO = cStringIO.StringIO
else:
    unichr = chr
    xrange = range
    string_types = (str,)
    text_type = str
    u_prefix = ''
    iteritems = dict.items
    itervalues = dict.values
    from io import StringIO, BytesIO

def add_metaclass(metaclass):
    """Class decorator for creating a class with a metaclass."""
    def wrapper(cls):
        orig_vars = cls.__dict__.copy()
        orig_vars.pop('__dict__', None)
        orig_vars.pop('__weakref__', None)
        for slots_var in orig_vars.get('__slots__', ()):
            orig_vars.pop(slots_var)
        return metaclass(cls.__name__, cls.__bases__, orig_vars)
    return wrapper
<commit_msg>Clarified comment in _surrogatepair() util.<commit_after># -*- coding: utf-8 -*-
"""
    pygments.util
    ~~~~~~~~~~~~~

    Utility functions.

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import re
import sys

split_path_re = re.compile(r'[/\\ ]')
doctype_lookup_re = re.compile(r'''(?smx)
    (<\?.*?\?>)?\s*
    <!DOCTYPE\s+(
     [a-zA-Z_][a-zA-Z0-9]*\s+
     [a-zA-Z_][a-zA-Z0-9]*\s+
     "[^"]*")
     [^>]*>
''')
tag_re = re.compile(r'<(.+?)(\s.*?)?>.*?</.+?>(?uism)')

class ClassNotFound(ValueError):
    """
    If one of the get_*_by_* functions didn't find a matching class.
    """

class OptionError(Exception):
    pass

def get_choice_opt(options, optname, allowed, default=None, normcase=False):
    string = options.get(optname, default)
    if normcase:
        string = string.lower()
    if string not in allowed:
        raise OptionError('Value for option %s must be one of %s' %
                          (optname, ', '.join(map(str, allowed))))
    return string

def get_bool_opt(options, optname, default=None):
    string = options.get(optname, default)
    if isinstance(string, bool):
        return string
    elif isinstance(string, int):
        return bool(string)
    elif not isinstance(string, string_types):
        raise OptionError('Invalid type %r for option %s; use '
                          '1/0, yes/no, true/false, on/off' % (
                          string, optname))
    elif string.lower() in ('1', 'yes', 'true', 'on'):
        return True
    elif string.lower() in ('0', 'no', 'false', 'off'):
        return False
    else:
        raise OptionError('Invalid value %r for option %s; use '
                          '1/0, yes/no, true/false, on/off' % (
                          string, optname))

def get_int_opt(options, optname, default=None):
    string = options.get(optname, default)
    try:
        return int(string)
    except TypeError:
        raise OptionError('Invalid type %r for option %s; you '
                          'must give an integer value' % (
                          string, optname))
    except ValueError:
        raise OptionError('Invalid value %r for option %s; you '
                          'must give an integer value' % (
                          string, optname))

def get_list_opt(options, optname, default=None):
    val = options.get(optname, default)
    if isinstance(val, string_types):
        return val.split()
    elif isinstance(val, (list, tuple)):
        return list(val)
    else:
        raise OptionError('Invalid type %r for option %s; you '
                          'must give a list value' % (
                          val, optname))

def docstring_headline(obj):
    if not obj.__doc__:
        return ''
    res = []
    for line in obj.__doc__.strip().splitlines():
        if line.strip():
            res.append(" " + line.strip())
        else:
            break
    return ''.join(res).lstrip()

def make_analysator(f):
    """
    Return a static text analysation function that
    returns float values.
    """
    def text_analyse(text):
        try:
            rv = f(text)
        except Exception:
            return 0.0
        if not rv:
            return 0.0
        try:
            return min(1.0, max(0.0, float(rv)))
        except (ValueError, TypeError):
            return 0.0
    text_analyse.__doc__ = f.__doc__
    return staticmethod(text_analyse)

def shebang_matches(text, regex):
    """
    Check if the given regular expression matches the last part of the
    shebang if one exists.

        >>> from pygments.util import shebang_matches
        >>> shebang_matches('#!/usr/bin/env python', r'python(2\.\d)?')
        True
        >>> shebang_matches('#!/usr/bin/python2.4', r'python(2\.\d)?')
        True
        >>> shebang_matches('#!/usr/bin/python-ruby', r'python(2\.\d)?')
        False
        >>> shebang_matches('#!/usr/bin/python/ruby', r'python(2\.\d)?')
        False
        >>> shebang_matches('#!/usr/bin/startsomethingwith python',
        ...                 r'python(2\.\d)?')
        True

    It also checks for common windows executable file extensions::

        >>> shebang_matches('#!C:\\Python2.4\\Python.exe', r'python(2\.\d)?')
        True

    Parameters (``'-f'`` or ``'--foo'`` are ignored so ``'perl'`` does
    the same as ``'perl -e'``)

    Note that this method automatically searches the whole string (eg:
    the regular expression is wrapped in ``'^$'``)
    """
    index = text.find('\n')
    if index >= 0:
        first_line = text[:index].lower()
    else:
        first_line = text.lower()
    if first_line.startswith('#!'):
        try:
            found = [x for x in split_path_re.split(first_line[2:].strip())
                     if x and not x.startswith('-')][-1]
        except IndexError:
            return False
        regex = re.compile('^%s(\.(exe|cmd|bat|bin))?$' % regex, re.IGNORECASE)
        if regex.search(found) is not None:
            return True
    return False

def doctype_matches(text, regex):
    """
    Check if the doctype matches a regular expression (if present).
    Note that this method only checks the first part of a DOCTYPE.
    eg: 'html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN"'
    """
    m = doctype_lookup_re.match(text)
    if m is None:
        return False
    doctype = m.group(2)
    return re.compile(regex).match(doctype.strip()) is not None

def html_doctype_matches(text):
    """
    Check if the file looks like it has a html doctype.
    """
    return doctype_matches(text, r'html\s+PUBLIC\s+"-//W3C//DTD X?HTML.*')

_looks_like_xml_cache = {}
def looks_like_xml(text):
    """
    Check if a doctype exists or if we have some tags.
    """
    key = hash(text)
    try:
        return _looks_like_xml_cache[key]
    except KeyError:
        m = doctype_lookup_re.match(text)
        if m is not None:
            return True
        rv = tag_re.search(text[:1000]) is not None
        _looks_like_xml_cache[key] = rv
        return rv

# Python narrow build compatibility

def _surrogatepair(c):
    # Given a unicode character code
    # with length greater than 16 bits,
    # return the two 16 bit surrogate pair.
    # From example D28 of:
    # http://www.unicode.org/book/ch03.pdf
    return (0xd7c0 + (c >> 10), (0xdc00 + (c & 0x3ff)))

def unirange(a, b):
    """
    Returns a regular expression string to match the given non-BMP range.
    """
    if b < a:
        raise ValueError("Bad character range")
    if a < 0x10000 or b < 0x10000:
        raise ValueError("unirange is only defined for non-BMP ranges")

    if sys.maxunicode > 0xffff:
        # wide build
        return u'[%s-%s]' % (unichr(a), unichr(b))
    else:
        # narrow build stores surrogates, and the 're' module handles them
        # (incorrectly) as characters.  Since there is still ordering among
        # these characters, expand the range to one that it understands.  Some
        # background in http://bugs.python.org/issue3665 and
        # http://bugs.python.org/issue12749
        #
        # Additionally, the lower constants are using unichr rather than
        # literals because jython [which uses the wide path] can't load this
        # file if they are literals.
        ah, al = _surrogatepair(a)
        bh, bl = _surrogatepair(b)
        if ah == bh:
            return u'(?:%s[%s-%s])' % (unichr(ah), unichr(al), unichr(bl))
        else:
            buf = []
            buf.append(u'%s[%s-%s]' %
                       (unichr(ah), unichr(al),
                        ah == bh and unichr(bl) or unichr(0xdfff)))
            if ah - bh > 1:
                buf.append(u'[%s-%s][%s-%s]' %
                           unichr(ah+1), unichr(bh-1), unichr(0xdc00), unichr(0xdfff))
            if ah != bh:
                buf.append(u'%s[%s-%s]' %
                           (unichr(bh), unichr(0xdc00), unichr(bl)))

            return u'(?:' + u'|'.join(buf) + u')'

# Python 2/3 compatibility

if sys.version_info < (3, 0):
    unichr = unichr
    xrange = xrange
    string_types = (str, unicode)
    text_type = unicode
    u_prefix = 'u'
    iteritems = dict.iteritems
    itervalues = dict.itervalues
    import StringIO, cStringIO
    # unfortunately, io.StringIO in Python 2 doesn't accept str at all
    StringIO = StringIO.StringIO
    BytesIO = cStringIO.StringIO
else:
    unichr = chr
    xrange = range
    string_types = (str,)
    text_type = str
    u_prefix = ''
    iteritems = dict.items
    itervalues = dict.values
    from io import StringIO, BytesIO

def add_metaclass(metaclass):
    """Class decorator for creating a class with a metaclass."""
    def wrapper(cls):
        orig_vars = cls.__dict__.copy()
        orig_vars.pop('__dict__', None)
        orig_vars.pop('__weakref__', None)
        for slots_var in orig_vars.get('__slots__', ()):
            orig_vars.pop(slots_var)
        return metaclass(cls.__name__, cls.__bases__, orig_vars)
    return wrapper
<|endoftext|>


<commit_before>from straight.plugin import load as plugin_load

class Format(object):
    @classmethod
    def name(cls):
        """Get the format's name.

        An format's canonical name is the lowercase name of the class
        implementing it.

        >>> Format().name()
        'format'
        """
        return cls.__name__.lower()

    def read(self, stream):
        """Read style from *stream* according to the format.
        """
        raise NotImplementedError

    def write(self, stream):
        """Write style to *stream* according to the format.
        """
        raise NotImplementedError

class PluginError(Exception):
    pass

def find_formats():
    plugins = plugin_load('pygout.formats', subclasses=Format)
    formats = {}

    for p in plugins:
        if p.name() in formats:
            raise PluginError(('Duplicate format: '
                '{e.__module__}.{e.__name__} and '
                '{n.__module__}.{n.__name__}').format(e=formats[p.name()],
                                                      n=p))
        else:
            formats[p.name()] = p

    return formats
<commit_msg>Fix stupid oops in Format interface<commit_after>from straight.plugin import load as plugin_load

class Format(object):
    @classmethod
    def name(cls):
        """Get the format's name.

        An format's canonical name is the lowercase name of the class
        implementing it.

        >>> Format().name()
        'format'
        """
        return cls.__name__.lower()

    def read(self, stream):
        """Read style from *stream* according to the format.
        """
        raise NotImplementedError

    def write(self, style, stream):
        """Write *style* to *stream* according to the format.
        """
        raise NotImplementedError

class PluginError(Exception):
    pass

def find_formats():
    plugins = plugin_load('pygout.formats', subclasses=Format)
    formats = {}

    for p in plugins:
        if p.name() in formats:
            raise PluginError(('Duplicate format: '
                '{e.__module__}.{e.__name__} and '
                '{n.__module__}.{n.__name__}').format(e=formats[p.name()],
                                                      n=p))
        else:
            formats[p.name()] = p

    return formats
<|endoftext|>


<commit_before># coding: utf8

import sys

from .components import ConfigObject, DependencyContainer, Router, View
from .components import initialize_loggers
from .http import Context
from .router_exceptions import (
    RouteSuccessfullyDispatched, RouteNotFound, RoutePassed, DispatchFinished
)
from .session import MockSession, SessionStorage, FileSessionStorage
from .session import RedisSessionStorage, RedisSentinelSessionStorage

from inspect import getmembers, ismethod
from json import dumps as json_dumps
from logging import getLogger
from os import path
from uwsgi import opt as uwsgi_opt

log = getLogger("pygrim.server")

class ResponseWrap(object):
    def __init__(self, start_response):
        self._start_response = start_response

    def __call__(self, status, headers):
        self._start_response(status, headers)
        self._start_response = self.noop

    def noop(self, status, headers):
        pass

class Server(object):

    KNOWN_CONFIG_FORMATS = (
        "yaml", "ini"
    )

    def __init__(self):
        self._initialize_basic_components()
        self._methods = {}
        self._not_found_method = None

    def __call__(self, environment, start_response):
        start_response = ResponseWrap(start_response)
        context = Context(environment)
        try:
            self._handle_request(context)
        except:
            log.exception("Fatal Error")
            start_response("500: Fatal Server Error", [])
            yield "Fatal Server Error"
        else:
            context.finalize_response()
            start_response(
                context.get_response_status_code(),
                context.get_response_headers()
            )
            yield context.get_response_body()

        return

    def display(self, *args, **kwargs):
        self._dic.view.display(*args, **kwargs)
        raise DispatchFinished()

    def redirect(self, context, **kwargs):
        if "url" in kwargs:
            context.redirect(kwargs.pop("url"), **kwargs)
        elif "route_name" in kwargs:
            url = "".join((
                context.get_request_url(),
                self._dic.router.url_for(
                    kwargs.pop("route_name"),
                    kwargs.pop("params", None) or {}
                )
            ))
            context.redirect(url, **kwargs)
        else:
            raise RuntimeError("redirect needs url or route_name params")
        raise DispatchFinished()

    # napojit na postfork
    def do_postfork(self):
        self._collect_exposed_methods()
        if hasattr(self, "_route_register_func"):
            self._route_register_func(self._dic.router)
            self._finalize_routes()
            log.debug("Routes loaded")
        else:
            log.warning("There is no function to register routes!")

        if hasattr(self, "postfork"):
            self.postfork()

    def render(self, *args, **kwargs):
        log.warning(
            "STOP CALLING AWFUL %r, CALL %r INSTEAD!",
            "server.render", "server.display"
        )
        return self.display(*args, **kwargs)

    def _collect_exposed_methods(self):
        def is_exposed(member, member_name):
            return (
                getattr(member, "_exposed", False) is True
            )

        for member_name, member in getmembers(self, predicate=ismethod):
            if is_exposed(member, member_name):
                try:
                    self._methods[member._dispatch_name] = member
                except AttributeError:
                    self._methods[member_name] = member

                if getattr(member, "_not_found", False) is True:
                    self._not_found_method = member

                if getattr(member, "_error", False) is True:
                    self._error_method = member

    def _default_error_method(self, context, exc):
        log.exception(exc.message)
        context.set_response_body("Internal Server Error")
        context.set_response_status(500)

    def _default_not_found_method(self, context):
        context.set_response_body("Not found")
        context.set_response_status(404)

    def _finalize_routes(self):
        for route in self._dic.router.get_routes():
            try:
                method = self._methods[route.get_handle_name()]
            except KeyError:
                raise RuntimeError(
                    u"Server has no method %r to handle route %r" % (
                        route._handle_name,
                        route.get_name() or route.get_pattern()
                    )
                )
            else:
                route.assign_method(method)

    def _find_session_handler(self):
        storage_type = self.config.get("session:type")
        if storage_type == "file":
            storage_class = FileSessionStorage
        elif storage_type == "redis":
            storage_class = RedisSessionStorage
        elif storage_type == "redis-sentinel":
            storage_class = RedisSentinelSessionStorage
        else:
            raise RuntimeError("Unknown session handler: %r", storage_type)

        return storage_class

    def _get_config_path(self):
        for key in self.KNOWN_CONFIG_FORMATS:
            if key in uwsgi_opt:
                return uwsgi_opt[key]
        else:
            raise RuntimeError("No known config format used to start uwsgi!")

    def _handle_request(self, context):
        try:
            session_loaded = False
            for route in self._dic.router.matching_routes(context):
                if route.requires_session() and context.session is None:
                    context.load_session(self.session_handler)
                    session_loaded = True
                    log.debug(
                        "Session handler:%r loaded session:%r",
                        type(self.session_handler), context.session
                    )

                try:
                    route.dispatch(context=context)
                except DispatchFinished:
                    pass
                except RoutePassed:
                    continue

                raise RouteSuccessfullyDispatched()
            else:
                not_found_method = (
                    self._not_found_method or self._default_not_found_method
                )
                try:
                    not_found_method(context=context)
                except DispatchFinished:
                    raise RouteNotFound
        except RouteSuccessfullyDispatched:
            log.debug("Dispatch succeded on: %r", context.current_route)
            if session_loaded:
                context.save_session(self.session_handler)

            return
        except RouteNotFound:
            return
        except:
            exc = sys.exc_info()
        log.error(
            "Error while dispatching to: %r",
            (
                context.current_route._handle_name
                if context.current_route
                else "<no route>"
            )
        )
        if hasattr(self, "_error_method"):
            try:
                self._error_method(context=context, exc=exc[1])
                return
            except DispatchFinished:
                return
            except:
                exc = sys.exc_info()[1]
        try:
            self._default_error_method(context=context, exc=exc[1])
        except:
            log.exception("Error in default_error_method")
            log.critical("Error in default_error_method")
            raise
        # endtry

    def _initialize_basic_components(self):
        self._dic = DependencyContainer()

        self.config = ConfigObject(self._get_config_path())
        self._register_logger(self.config)
        self._dic.mode = self.config.get("grim:mode")
        self._dic.router = Router()

        if self.config.get("view:enabled", True):
            self._register_view(self.config)

        self._register_session_handler()
        log.debug("Basic components initialized")

    def _register_session_handler(self):
        if self.config.get("session:enabled", False):
            storage_class = self._find_session_handler()
        else:
            storage_class = MockSession

        handler = storage_class(self.config)
        if not isinstance(handler, SessionStorage):
            raise ValueError(
                "SessionHandler should be derived from SessionStorage"
            )

        self.session_handler = handler

    def _register_logger(self, config):
        initialize_loggers(config)

    def _register_view(self, config):
        extra_functions = {
            "base_url": self._jinja_base_url,
            "site_url": self._jinja_site_url,
            "url_for": self._jinja_url_for,
            "as_json": self._jinja_as_json,
        }
        self._dic.view = View(config, extra_functions)

    # jinja extra methods
    def _jinja_as_json(self, data):
        return json_dumps(data)

    def _jinja_base_url(self, context):
        return "%s%s" % (
            context.get_request_url(), context.get_request_root_uri()
        )

    def _jinja_site_url(self, context, site):
        return path.join(self._jinja_base_url(context), site)

    def _jinja_url_for(self, route, params=None):
        params = params or {}
        url = self._dic.router.url_for(route, params)
        return url

# eof
<commit_msg>fix error metody<commit_after># coding: utf8

import sys

from .components import ConfigObject, DependencyContainer, Router, View
from .components import initialize_loggers
from .http import Context
from .router_exceptions import (
    RouteSuccessfullyDispatched, RouteNotFound, RoutePassed, DispatchFinished
)
from .session import MockSession, SessionStorage, FileSessionStorage
from .session import RedisSessionStorage, RedisSentinelSessionStorage

from inspect import getmembers, ismethod
from json import dumps as json_dumps
from logging import getLogger
from os import path
from uwsgi import opt as uwsgi_opt

log = getLogger("pygrim.server")

class ResponseWrap(object):
    def __init__(self, start_response):
        self._start_response = start_response

    def __call__(self, status, headers):
        self._start_response(status, headers)
        self._start_response = self.noop

    def noop(self, status, headers):
        pass

class Server(object):

    KNOWN_CONFIG_FORMATS = (
        "yaml", "ini"
    )

    def __init__(self):
        self._initialize_basic_components()
        self._methods = {}
        self._not_found_method = None

    def __call__(self, environment, start_response):
        start_response = ResponseWrap(start_response)
        context = Context(environment)
        try:
            self._handle_request(context)
        except:
            log.exception("Fatal Error")
            start_response("500: Fatal Server Error", [])
            yield "Fatal Server Error"
        else:
            context.finalize_response()
            start_response(
                context.get_response_status_code(),
                context.get_response_headers()
            )
            yield context.get_response_body()

        return

    def display(self, *args, **kwargs):
        self._dic.view.display(*args, **kwargs)
        raise DispatchFinished()

    def redirect(self, context, **kwargs):
        if "url" in kwargs:
            context.redirect(kwargs.pop("url"), **kwargs)
        elif "route_name" in kwargs:
            url = "".join((
                context.get_request_url(),
                self._dic.router.url_for(
                    kwargs.pop("route_name"),
                    kwargs.pop("params", None) or {}
                )
            ))
            context.redirect(url, **kwargs)
        else:
            raise RuntimeError("redirect needs url or route_name params")
        raise DispatchFinished()

    # napojit na postfork
    def do_postfork(self):
        self._collect_exposed_methods()
        if hasattr(self, "_route_register_func"):
            self._route_register_func(self._dic.router)
            self._finalize_routes()
            log.debug("Routes loaded")
        else:
            log.warning("There is no function to register routes!")

        if hasattr(self, "postfork"):
            self.postfork()

    def render(self, *args, **kwargs):
        log.warning(
            "STOP CALLING AWFUL %r, CALL %r INSTEAD!",
            "server.render", "server.display"
        )
        return self.display(*args, **kwargs)

    def _collect_exposed_methods(self):
        def is_exposed(member, member_name):
            return (
                getattr(member, "_exposed", False) is True
            )

        for member_name, member in getmembers(self, predicate=ismethod):
            if is_exposed(member, member_name):
                try:
                    self._methods[member._dispatch_name] = member
                except AttributeError:
                    self._methods[member_name] = member

                if getattr(member, "_not_found", False) is True:
                    self._not_found_method = member

                if getattr(member, "_error", False) is True:
                    self._error_method = member

    def _default_error_method(self, context, exc):
        log.exception(exc.message)
        context.set_response_body("Internal Server Error")
        context.set_response_status(500)

    def _default_not_found_method(self, context):
        context.set_response_body("Not found")
        context.set_response_status(404)

    def _finalize_routes(self):
        for route in self._dic.router.get_routes():
            try:
                method = self._methods[route.get_handle_name()]
            except KeyError:
                raise RuntimeError(
                    u"Server has no method %r to handle route %r" % (
                        route._handle_name,
                        route.get_name() or route.get_pattern()
                    )
                )
            else:
                route.assign_method(method)

    def _find_session_handler(self):
        storage_type = self.config.get("session:type")
        if storage_type == "file":
            storage_class = FileSessionStorage
        elif storage_type == "redis":
            storage_class = RedisSessionStorage
        elif storage_type == "redis-sentinel":
            storage_class = RedisSentinelSessionStorage
        else:
            raise RuntimeError("Unknown session handler: %r", storage_type)

        return storage_class

    def _get_config_path(self):
        for key in self.KNOWN_CONFIG_FORMATS:
            if key in uwsgi_opt:
                return uwsgi_opt[key]
        else:
            raise RuntimeError("No known config format used to start uwsgi!")

    def _handle_request(self, context):
        try:
            session_loaded = False
            for route in self._dic.router.matching_routes(context):
                if route.requires_session() and context.session is None:
                    context.load_session(self.session_handler)
                    session_loaded = True
                    log.debug(
                        "Session handler:%r loaded session:%r",
                        type(self.session_handler), context.session
                    )

                try:
                    route.dispatch(context=context)
                except DispatchFinished:
                    pass
                except RoutePassed:
                    continue

                raise RouteSuccessfullyDispatched()
            else:
                not_found_method = (
                    self._not_found_method or self._default_not_found_method
                )
                try:
                    not_found_method(context=context)
                except DispatchFinished:
                    raise RouteNotFound
        except RouteSuccessfullyDispatched:
            log.debug("Dispatch succeded on: %r", context.current_route)
            if session_loaded:
                context.save_session(self.session_handler)

            return
        except RouteNotFound:
            return
        except:
            exc = sys.exc_info()[1]
        log.error(
            "Error while dispatching to: %r",
            (
                context.current_route._handle_name
                if context.current_route
                else "<no route>"
            )
        )
        if hasattr(self, "_error_method"):
            try:
                self._error_method(context=context, exc=exc)
                return
            except DispatchFinished:
                return
            except:
                exc = sys.exc_info()[1]
        try:
            self._default_error_method(context=context, exc=exc)
        except:
            log.exception("Error in default_error_method")
            log.critical("Error in default_error_method")
            raise
        # endtry

    def _initialize_basic_components(self):
        self._dic = DependencyContainer()

        self.config = ConfigObject(self._get_config_path())
        self._register_logger(self.config)
        self._dic.mode = self.config.get("grim:mode")
        self._dic.router = Router()

        if self.config.get("view:enabled", True):
            self._register_view(self.config)

        self._register_session_handler()
        log.debug("Basic components initialized")

    def _register_session_handler(self):
        if self.config.get("session:enabled", False):
            storage_class = self._find_session_handler()
        else:
            storage_class = MockSession

        handler = storage_class(self.config)
        if not isinstance(handler, SessionStorage):
            raise ValueError(
                "SessionHandler should be derived from SessionStorage"
            )

        self.session_handler = handler

    def _register_logger(self, config):
        initialize_loggers(config)

    def _register_view(self, config):
        extra_functions = {
            "base_url": self._jinja_base_url,
            "site_url": self._jinja_site_url,
            "url_for": self._jinja_url_for,
            "as_json": self._jinja_as_json,
        }
        self._dic.view = View(config, extra_functions)

    # jinja extra methods
    def _jinja_as_json(self, data):
        return json_dumps(data)

    def _jinja_base_url(self, context):
        return "%s%s" % (
            context.get_request_url(), context.get_request_root_uri()
        )

    def _jinja_site_url(self, context, site):
        return path.join(self._jinja_base_url(context), site)

    def _jinja_url_for(self, route, params=None):
        params = params or {}
        url = self._dic.router.url_for(route, params)
        return url

# eof
<|endoftext|>


<commit_before># coding: utf8

import sys

from .components import ConfigObject, DependencyContainer, Router, View
from .components import initialize_loggers
from .http import Context
from .router_exceptions import (
    RouteSuccessfullyDispatched, RouteNotFound, RoutePassed
)
from .session import MockSession, SessionStorage, FileSessionStorage

from inspect import getmembers, ismethod
from json import dumps as json_dumps
from logging import getLogger
from os import path
from uwsgi import opt as uwsgi_opt

log = getLogger("pygrim.server")

class ResponseWrap(object):
    def __init__(self, start_response):
        self._start_response = start_response

    def __call__(self, status, headers):
        self._start_response(status, headers)
        self._start_response = self.noop

    def noop(self, status, headers):
        pass

class Server(object):

    KNOWN_CONFIG_FORMATS = (
        "yaml", "ini"
    )

    def __init__(self):
        self._initialize_basic_components()
        self._methods = {}
        self._not_found_method = None

    def __call__(self, environment, start_response):
        start_response = ResponseWrap(start_response)
        context = Context(environment)
        self._handle_request(context)
        context.finalize_response()
        start_response(
            context.get_response_status_code(), context.get_response_headers()
        )
        yield context.get_response_body()
        return

    def display(self, *args, **kwargs):
        self._dic.view.display(*args, **kwargs)
        raise RouteSuccessfullyDispatched()

    def redirect(context, url, **kwargs):
        context.redirect(url, **kwargs)
        raise RouteSuccessfullyDispatched()

    # napojit na postfork
    def do_postfork(self):
        self._collect_exposed_methods()
        if hasattr(self, "_route_register_func"):
            self._route_register_func(self._dic.router)
            self._finalize_routes()
            log.debug("Routes loaded")
        else:
            log.warning("There is no function to register routes!")

        if hasattr(self, "postfork"):
            self.postfork()

    def render(self, *args, **kwargs):
        log.warning(
            "STOP CALLING AWFUL %r, CALL %r INSTEAD!",
            "server.render", "server.display"
        )
        return self.display(*args, **kwargs)

    def _collect_exposed_methods(self):
        def is_exposed(member, member_name):
            return (
                getattr(member, "_exposed", False) is True
            )

        for member_name, member in getmembers(self, predicate=ismethod):
            if is_exposed(member, member_name):
                try:
                    self._methods[member._dispatch_name] = member
                except AttributeError:
                    self._methods[member_name] = member

                if getattr(member, "_not_found", False) is True:
                    self._not_found_method = member

                if getattr(member, "_error", False) is True:
                    self._error_method = member

    def _default_error_method(self, context, exc):
        log.exception(exc.message)
        context.set_response_body("Internal Server Error")
        context.set_response_status(500)

    def _default_not_found_method(self, context):
        context.set_response_body("Not found")
        context.set_response_status(404)

    def _finalize_routes(self):
        for route in self._dic.router.get_routes():
            try:
                method = self._methods[route.get_handle_name()]
            except KeyError:
                raise RuntimeError(
                    u"Server has no method %r to handle route %r" % (
                        route._handle_name,
                        route.get_name() or route.get_pattern()
                    )
                )
            else:
                route.assign_method(method)

    def _find_session_handler(self):
        storage_type = self.config.get("session:type")
        if storage_type == "file":
            storage_class = FileSessionStorage
        else:
            raise RuntimeError("Unknown session handler: %r", storage_type)

        return storage_class

    def _get_config_path(self):
        for key in self.KNOWN_CONFIG_FORMATS:
            if key in uwsgi_opt:
                return uwsgi_opt[key]
        else:
            raise RuntimeError("No known config format used to start uwsgi!")

    def _handle_request(self, context):
        try:
            for route in self._dic.router.matching_routes(context):
                if route.requires_session() and context.session is None:
                    context.load_session(self.session_handler)
                    log.debug(
                        "Session handler:%r loaded session:%r",
                        type(self.session_handler), context.session
                    )

                try:
                    route.dispatch(context=context)
                except RouteSuccessfullyDispatched:
                    pass
                except RoutePassed:
                    pass
                if route.requires_session():
                    context.save_session(self.session_handler)

                raise RouteSuccessfullyDispatched()
            else:
                if self._not_found_method:
                    self._not_found_method(context=context)
                else:
                    raise RouteNotFound()
        except RouteSuccessfullyDispatched:
            log.debug("Dispatch succeded on: %r", context.current_route)
        except RouteNotFound:
            try:
                self._default_not_found_method(context=context)
            except RouteSuccessfullyDispatched:
                pass
        except:
            log.error(
                "Error while dispatching to: %r",
                (
                    context.current_route._handle_name
                    if context.current_route
                    else "<no route>"
                )
            )
            exc = sys.exc_info()[1]
            if hasattr(self, "_error_method"):
                try:
                    self._error_method(context=context, exc=exc)
                    return
                except RouteSuccessfullyDispatched:
                    return
                except:
                    exc = sys.exc_info()[1]

            self._default_error_method(context=context, exc=exc)

    def _initialize_basic_components(self):
        self._dic = DependencyContainer()

        self.config = ConfigObject(self._get_config_path())
        self._register_logger(self.config)
        self._dic.mode = self.config.get("grim:mode")
        self._dic.router = Router()

        if self.config.get("view:enabled", True):
            self._register_view(self.config)

        self._register_session_handler()
        log.debug("Basic components initialized")

    def _register_session_handler(self):
        if self.config.get("session:enabled", False):
            storage_class = self._find_session_handler()
        else:
            storage_class = MockSession

        handler = storage_class(self.config)
        if not isinstance(handler, SessionStorage):
            raise ValueError(
                "SessionHandler should be derived from SessionStorage"
            )

        self.session_handler = handler

    def _register_logger(self, config):
        initialize_loggers(config)

    def _register_view(self, config):
        extra_functions = {
            "base_url": self._jinja_base_url,
            "site_url": self._jinja_site_url,
            "url_for": self._jinja_url_for,
            "as_json": self._jinja_as_json,
        }
        self._dic.view = View(config, extra_functions)

    # jinja extra methods
    def _jinja_as_json(self, data):
        return json_dumps(data)

    def _jinja_base_url(self, context):
        return "%s%s" % (
            context.get_request_url(), context.get_request_root_uri()
        )

    def _jinja_site_url(self, context, site):
        return path.join(self._jinja_base_url(context), site)

    def _jinja_url_for(self, route, params=None):
        params = params or {}
        url = self._dic.router.url_for(route, params)
        return url

# eof
<commit_msg>redirect now can get route_name<commit_after># coding: utf8

import sys

from .components import ConfigObject, DependencyContainer, Router, View
from .components import initialize_loggers
from .http import Context
from .router_exceptions import (
    RouteSuccessfullyDispatched, RouteNotFound, RoutePassed
)
from .session import MockSession, SessionStorage, FileSessionStorage

from inspect import getmembers, ismethod
from json import dumps as json_dumps
from logging import getLogger
from os import path
from uwsgi import opt as uwsgi_opt

log = getLogger("pygrim.server")

class ResponseWrap(object):
    def __init__(self, start_response):
        self._start_response = start_response

    def __call__(self, status, headers):
        self._start_response(status, headers)
        self._start_response = self.noop

    def noop(self, status, headers):
        pass

class Server(object):

    KNOWN_CONFIG_FORMATS = (
        "yaml", "ini"
    )

    def __init__(self):
        self._initialize_basic_components()
        self._methods = {}
        self._not_found_method = None

    def __call__(self, environment, start_response):
        start_response = ResponseWrap(start_response)
        context = Context(environment)
        self._handle_request(context)
        context.finalize_response()
        start_response(
            context.get_response_status_code(), context.get_response_headers()
        )
        yield context.get_response_body()
        return

    def display(self, *args, **kwargs):
        self._dic.view.display(*args, **kwargs)
        raise RouteSuccessfullyDispatched()

    def redirect(self, context, **kwargs):
        if "url" in kwargs:
            context.redirect(kwargs.pop("url"), **kwargs)
        elif "route_name" in kwargs:
            url = "".join((
                context.get_request_url(),
                self._dic.router.url_for(
                    kwargs.pop("route_name"),
                    kwargs.pop("params", None) or {}
                )
            ))
            context.redirect(url, **kwargs)
        else:
            raise RuntimeError("redirect needs url or route_name params")
        raise RouteSuccessfullyDispatched()

    # napojit na postfork
    def do_postfork(self):
        self._collect_exposed_methods()
        if hasattr(self, "_route_register_func"):
            self._route_register_func(self._dic.router)
            self._finalize_routes()
            log.debug("Routes loaded")
        else:
            log.warning("There is no function to register routes!")

        if hasattr(self, "postfork"):
            self.postfork()

    def render(self, *args, **kwargs):
        log.warning(
            "STOP CALLING AWFUL %r, CALL %r INSTEAD!",
            "server.render", "server.display"
        )
        return self.display(*args, **kwargs)

    def _collect_exposed_methods(self):
        def is_exposed(member, member_name):
            return (
                getattr(member, "_exposed", False) is True
            )

        for member_name, member in getmembers(self, predicate=ismethod):
            if is_exposed(member, member_name):
                try:
                    self._methods[member._dispatch_name] = member
                except AttributeError:
                    self._methods[member_name] = member

                if getattr(member, "_not_found", False) is True:
                    self._not_found_method = member

                if getattr(member, "_error", False) is True:
                    self._error_method = member

    def _default_error_method(self, context, exc):
        log.exception(exc.message)
        context.set_response_body("Internal Server Error")
        context.set_response_status(500)

    def _default_not_found_method(self, context):
        context.set_response_body("Not found")
        context.set_response_status(404)

    def _finalize_routes(self):
        for route in self._dic.router.get_routes():
            try:
                method = self._methods[route.get_handle_name()]
            except KeyError:
                raise RuntimeError(
                    u"Server has no method %r to handle route %r" % (
                        route._handle_name,
                        route.get_name() or route.get_pattern()
                    )
                )
            else:
                route.assign_method(method)

    def _find_session_handler(self):
        storage_type = self.config.get("session:type")
        if storage_type == "file":
            storage_class = FileSessionStorage
        else:
            raise RuntimeError("Unknown session handler: %r", storage_type)

        return storage_class

    def _get_config_path(self):
        for key in self.KNOWN_CONFIG_FORMATS:
            if key in uwsgi_opt:
                return uwsgi_opt[key]
        else:
            raise RuntimeError("No known config format used to start uwsgi!")

    def _handle_request(self, context):
        try:
            for route in self._dic.router.matching_routes(context):
                if route.requires_session() and context.session is None:
                    context.load_session(self.session_handler)
                    log.debug(
                        "Session handler:%r loaded session:%r",
                        type(self.session_handler), context.session
                    )

                try:
                    route.dispatch(context=context)
                except RouteSuccessfullyDispatched:
                    pass
                except RoutePassed:
                    pass
                if route.requires_session():
                    context.save_session(self.session_handler)

                raise RouteSuccessfullyDispatched()
            else:
                if self._not_found_method:
                    self._not_found_method(context=context)
                else:
                    raise RouteNotFound()
        except RouteSuccessfullyDispatched:
            log.debug("Dispatch succeded on: %r", context.current_route)
        except RouteNotFound:
            try:
                self._default_not_found_method(context=context)
            except RouteSuccessfullyDispatched:
                pass
        except:
            log.error(
                "Error while dispatching to: %r",
                (
                    context.current_route._handle_name
                    if context.current_route
                    else "<no route>"
                )
            )
            exc = sys.exc_info()[1]
            if hasattr(self, "_error_method"):
                try:
                    self._error_method(context=context, exc=exc)
                    return
                except RouteSuccessfullyDispatched:
                    return
                except:
                    exc = sys.exc_info()[1]

            self._default_error_method(context=context, exc=exc)

    def _initialize_basic_components(self):
        self._dic = DependencyContainer()

        self.config = ConfigObject(self._get_config_path())
        self._register_logger(self.config)
        self._dic.mode = self.config.get("grim:mode")
        self._dic.router = Router()

        if self.config.get("view:enabled", True):
            self._register_view(self.config)

        self._register_session_handler()
        log.debug("Basic components initialized")

    def _register_session_handler(self):
        if self.config.get("session:enabled", False):
            storage_class = self._find_session_handler()
        else:
            storage_class = MockSession

        handler = storage_class(self.config)
        if not isinstance(handler, SessionStorage):
            raise ValueError(
                "SessionHandler should be derived from SessionStorage"
            )

        self.session_handler = handler

    def _register_logger(self, config):
        initialize_loggers(config)

    def _register_view(self, config):
        extra_functions = {
            "base_url": self._jinja_base_url,
            "site_url": self._jinja_site_url,
            "url_for": self._jinja_url_for,
            "as_json": self._jinja_as_json,
        }
        self._dic.view = View(config, extra_functions)

    # jinja extra methods
    def _jinja_as_json(self, data):
        return json_dumps(data)

    def _jinja_base_url(self, context):
        return "%s%s" % (
            context.get_request_url(), context.get_request_root_uri()
        )

    def _jinja_site_url(self, context, site):
        return path.join(self._jinja_base_url(context), site)

    def _jinja_url_for(self, route, params=None):
        params = params or {}
        url = self._dic.router.url_for(route, params)
        return url

# eof
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

__title__ = 'pyhn'
__version__ = '0.3.0'
__author__ = 'Geoffrey Lehée'
__license__ = 'MIT'
__copyright__ = 'Copyright 2015 Geoffrey Lehée'
<commit_msg>0.3.2<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

__title__ = 'pyhn'
__version__ = '0.3.2'
__author__ = 'Geoffrey Lehée'
__license__ = 'MIT'
__copyright__ = 'Copyright 2015 Geoffrey Lehée'
<|endoftext|>