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

from pircel import model, tornado_adapter

logger = logging.getLogger(__name__)

commands = {'join',
            'part',
            'query',
            'me',
            'nick',
            'connect',
            'help',
            }

def build_prefix_map(strings):
    out = collections.defaultdict(list)
    for string in strings:
        prefix = ''
        for char in string:
            prefix += char
            out[prefix].append(string)
    return out

prefix_commands = build_prefix_map(commands)

class CommandParser(argparse.ArgumentParser):
    def parse_args(self, buffer, args_or_empty_list, *args, **kwargs):
        self.buffer = buffer
        try:
            split_args = args_or_empty_list[0].split()
        except IndexError:
            split_args = []
        return super(CommandParser, self).parse_args(split_args, *args, **kwargs)

    def _print_message(self, message, unused_file=None):
        buffer = self.buffer
        model.create_line(buffer=buffer, content='=' * 80, kind='other', nick='-*-')
        for line in message.splitlines():
            model.create_line(buffer=buffer, content=line, kind='other', nick='-*-')
        model.create_line(buffer=buffer, content='=' * 80, kind='other', nick='-*-')

    def exit(self, *args, **kwargs):
        raise ValueError()

    def decorate(self, function):
        def inner_function(instance, buffer, args_or_empty_list):
            try:
                args = self.parse_args(buffer, args_or_empty_list)
            except ValueError:
                pass
            else:
                args.buffer = buffer
                function(instance, args)

        inner_function.parser = self
        return inner_function

help_parser = CommandParser(prog='help', description='Display help and usage information for commands')
help_parser.add_argument('command', help='The command to display help for', choices=commands)

join_parser = CommandParser(prog='join', description='Join a new channel')
join_parser.add_argument('channel', help='The channel to join')
join_parser.add_argument('password', default=None, nargs='?',
                         help='Optional password for the channel')

part_parser = CommandParser(prog='part', description='Leave a channel')
part_parser.add_argument('channel', help='The channel to leave', default=None, nargs='?')

query_parser = CommandParser(prog='query', description='Start a private conversation')
query_parser.add_argument('who', help='Who to start a conversation with')

nick_parser = CommandParser(prog='query', description='Change your nickname on this server')
nick_parser.add_argument('new_nick', help='What to change your nick to')

me_parser = CommandParser(prog='me', description='Do a thing!')
me_parser.add_argument('action', help='The thing to do', nargs=argparse.REMAINDER)

connect_parser = CommandParser(prog='connect', description='Connect to a new IRC server')
connect_parser.add_argument('-i', '--insecure', action='store_true',
                            help='Enable ssl/tls for this server')
connect_parser.add_argument('-p', '--port', default=6697,
                            help='The port to connect on')
connect_parser.add_argument('-n', '--nick', default=None,
                            help='The nick to use on this server')
connect_parser.add_argument('-r', '--realname', default=None,
                            help='The real name to use on this server')
connect_parser.add_argument('-u', '--username', default=None,
                            help='The username to use on this server')
connect_parser.add_argument('host', help='The server to connect to')

class Dispatcher:
    def __init__(self, interfaces):
        self.interfaces = interfaces

    def dispatch(self, buffer_id, line):
        if line.startswith('/'):
            line = line[1:]
        command, *rest = line.split(maxsplit=1)
        command = command.lower()

        if command in prefix_commands:
            buffer = model.IRCBufferModel.get(id=buffer_id)
            if len(prefix_commands[command]) == 1:
                actual_command, = prefix_commands[command]
                getattr(self, actual_command)(buffer, rest)
            else:
                model.create_line(buffer=buffer,
                                  content='ambiguous command "{}"'.format(command),
                                  kind='other',
                                  nick='-*-')

    @help_parser.decorate
    def help(self, args):
        parser = getattr(self, args.command).parser
        parser.buffer = args.buffer
        parser.print_help()

    @join_parser.decorate
    def join(self, args):
        if args.channel is None:
            args.channel = args.buffer.name
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.join(args.channel, args.password)

    @part_parser.decorate
    def part(self, args):
        if args.channel is None:
            args.channel = args.buffer.name
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.part(args.channel)

    @query_parser.decorate
    def query(self, args):
        model.ensure_buffer(args.who, args.buffer.server)

    def me(self, buffer, rest):
        line = rest[0]
        interface = self.interfaces[buffer.server.id]
        interface.server_handler.send_message(buffer.name, '\1ACTION {}\1'.format(line))

    # Doesn't actually use the parser but we want /help to work
    me.parser = me_parser

    @nick_parser.decorate
    def nick(self, args):
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.change_nick(args.new_nick)

    @connect_parser.decorate
    def connect(self, args):
        if None in {args.nick, args.realname, args.username}:
            user = args.buffer.server.user
        server = model.create_server(host=args.host,
                                     port=args.port,
                                     secure=not args.insecure,
                                     nick=args.nick or user.nick,
                                     realname=args.realname or user.realname,
                                     username=args.username or user.username)

        interface = model.IRCServerInterface(server)
        tornado_adapter.IRCClient.from_interface(interface).connect()
        self.interfaces[interface.server_model.id] = interface

def main():
    pass

if __name__ == '__main__':
    main()
<commit_msg>Fix help text for /connect with --insecure<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import collections
import logging

from pircel import model, tornado_adapter

logger = logging.getLogger(__name__)

commands = {'join',
            'part',
            'query',
            'me',
            'nick',
            'connect',
            'help',
            }

def build_prefix_map(strings):
    out = collections.defaultdict(list)
    for string in strings:
        prefix = ''
        for char in string:
            prefix += char
            out[prefix].append(string)
    return out

prefix_commands = build_prefix_map(commands)

class CommandParser(argparse.ArgumentParser):
    def parse_args(self, buffer, args_or_empty_list, *args, **kwargs):
        self.buffer = buffer
        try:
            split_args = args_or_empty_list[0].split()
        except IndexError:
            split_args = []
        return super(CommandParser, self).parse_args(split_args, *args, **kwargs)

    def _print_message(self, message, unused_file=None):
        buffer = self.buffer
        model.create_line(buffer=buffer, content='=' * 80, kind='other', nick='-*-')
        for line in message.splitlines():
            model.create_line(buffer=buffer, content=line, kind='other', nick='-*-')
        model.create_line(buffer=buffer, content='=' * 80, kind='other', nick='-*-')

    def exit(self, *args, **kwargs):
        raise ValueError()

    def decorate(self, function):
        def inner_function(instance, buffer, args_or_empty_list):
            try:
                args = self.parse_args(buffer, args_or_empty_list)
            except ValueError:
                pass
            else:
                args.buffer = buffer
                function(instance, args)

        inner_function.parser = self
        return inner_function

help_parser = CommandParser(prog='help', description='Display help and usage information for commands')
help_parser.add_argument('command', help='The command to display help for', choices=commands)

join_parser = CommandParser(prog='join', description='Join a new channel')
join_parser.add_argument('channel', help='The channel to join')
join_parser.add_argument('password', default=None, nargs='?',
                         help='Optional password for the channel')

part_parser = CommandParser(prog='part', description='Leave a channel')
part_parser.add_argument('channel', help='The channel to leave', default=None, nargs='?')

query_parser = CommandParser(prog='query', description='Start a private conversation')
query_parser.add_argument('who', help='Who to start a conversation with')

nick_parser = CommandParser(prog='query', description='Change your nickname on this server')
nick_parser.add_argument('new_nick', help='What to change your nick to')

me_parser = CommandParser(prog='me', description='Do a thing!')
me_parser.add_argument('action', help='The thing to do', nargs=argparse.REMAINDER)

connect_parser = CommandParser(prog='connect', description='Connect to a new IRC server')
connect_parser.add_argument('-i', '--insecure', action='store_true',
                            help='Disable ssl/tls for this server')
connect_parser.add_argument('-p', '--port', default=6697,
                            help='The port to connect on')
connect_parser.add_argument('-n', '--nick', default=None,
                            help='The nick to use on this server')
connect_parser.add_argument('-r', '--realname', default=None,
                            help='The real name to use on this server')
connect_parser.add_argument('-u', '--username', default=None,
                            help='The username to use on this server')
connect_parser.add_argument('host', help='The server to connect to')

class Dispatcher:
    def __init__(self, interfaces):
        self.interfaces = interfaces

    def dispatch(self, buffer_id, line):
        if line.startswith('/'):
            line = line[1:]
        command, *rest = line.split(maxsplit=1)
        command = command.lower()

        if command in prefix_commands:
            buffer = model.IRCBufferModel.get(id=buffer_id)
            if len(prefix_commands[command]) == 1:
                actual_command, = prefix_commands[command]
                getattr(self, actual_command)(buffer, rest)
            else:
                model.create_line(buffer=buffer,
                                  content='ambiguous command "{}"'.format(command),
                                  kind='other',
                                  nick='-*-')

    @help_parser.decorate
    def help(self, args):
        parser = getattr(self, args.command).parser
        parser.buffer = args.buffer
        parser.print_help()

    @join_parser.decorate
    def join(self, args):
        if args.channel is None:
            args.channel = args.buffer.name
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.join(args.channel, args.password)

    @part_parser.decorate
    def part(self, args):
        if args.channel is None:
            args.channel = args.buffer.name
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.part(args.channel)

    @query_parser.decorate
    def query(self, args):
        model.ensure_buffer(args.who, args.buffer.server)

    def me(self, buffer, rest):
        line = rest[0]
        interface = self.interfaces[buffer.server.id]
        interface.server_handler.send_message(buffer.name, '\1ACTION {}\1'.format(line))

    # Doesn't actually use the parser but we want /help to work
    me.parser = me_parser

    @nick_parser.decorate
    def nick(self, args):
        interface = self.interfaces[args.buffer.server.id]
        interface.server_handler.change_nick(args.new_nick)

    @connect_parser.decorate
    def connect(self, args):
        if None in {args.nick, args.realname, args.username}:
            user = args.buffer.server.user
        server = model.create_server(host=args.host,
                                     port=args.port,
                                     secure=not args.insecure,
                                     nick=args.nick or user.nick,
                                     realname=args.realname or user.realname,
                                     username=args.username or user.username)

        interface = model.IRCServerInterface(server)
        tornado_adapter.IRCClient.from_interface(interface).connect()
        self.interfaces[interface.server_model.id] = interface

def main():
    pass

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


<commit_before># Copyright 2012 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

"""Top-level presubmit script for swarm_client.

See http://dev.chromium.org/developers/how-tos/depottools/presubmit-scripts for
details on the presubmit API built into gcl.
"""

def CommonChecks(input_api, output_api):
  import sys
  def join(*args):
    return input_api.os_path.join(input_api.PresubmitLocalPath(), *args)

  output = []
  sys_path_backup = sys.path
  try:
    sys.path = [
      input_api.PresubmitLocalPath(),
      join('tests'),
      join('third_party'),
    ] + sys.path
    black_list = list(input_api.DEFAULT_BLACK_LIST) + [
        r'.*_pb2\.py$',
        r'.*_pb2_grpc\.py$',
    ]
    output.extend(input_api.canned_checks.RunPylint(
        input_api, output_api,
    black_list=black_list))
  finally:
    sys.path = sys_path_backup

  # These tests are touching the live infrastructure. It's a pain if your IP
  # is not whitelisted so do not run them for now. They should use a local fake
  # web service instead.
  blacklist = [
    r'.*isolateserver_smoke_test\.py$',
    r'.*isolateserver_load_test\.py$',
    r'.*swarming_smoke_test\.py$',
  ]
  unit_tests = input_api.canned_checks.GetUnitTestsRecursively(
      input_api, output_api,
      input_api.os_path.join(input_api.PresubmitLocalPath()),
      whitelist=[r'.+_test\.py$'],
      blacklist=blacklist)
  output.extend(input_api.RunTests(unit_tests))
  return output

# pylint: disable=unused-argument
def CheckChangeOnUpload(input_api, output_api):
  return []

def CheckChangeOnCommit(input_api, output_api):
  return CommonChecks(input_api, output_api)
<commit_msg>[swarming] delete old blacklist file pattern<commit_after># Copyright 2012 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

"""Top-level presubmit script for swarm_client.

See http://dev.chromium.org/developers/how-tos/depottools/presubmit-scripts for
details on the presubmit API built into gcl.
"""

def CommonChecks(input_api, output_api):
  import sys
  def join(*args):
    return input_api.os_path.join(input_api.PresubmitLocalPath(), *args)

  output = []
  sys_path_backup = sys.path
  try:
    sys.path = [
      input_api.PresubmitLocalPath(),
      join('tests'),
      join('third_party'),
    ] + sys.path
    black_list = list(input_api.DEFAULT_BLACK_LIST) + [
        r'.*_pb2\.py$',
        r'.*_pb2_grpc\.py$',
    ]
    output.extend(input_api.canned_checks.RunPylint(
        input_api, output_api,
    black_list=black_list))
  finally:
    sys.path = sys_path_backup

  # These tests are touching the live infrastructure. It's a pain if your IP
  # is not whitelisted so do not run them for now. They should use a local fake
  # web service instead.
  blacklist = [
    r'.*isolateserver_smoke_test\.py$',
    r'.*isolateserver_load_test\.py$',
  ]
  unit_tests = input_api.canned_checks.GetUnitTestsRecursively(
      input_api, output_api,
      input_api.os_path.join(input_api.PresubmitLocalPath()),
      whitelist=[r'.+_test\.py$'],
      blacklist=blacklist)
  output.extend(input_api.RunTests(unit_tests))
  return output

# pylint: disable=unused-argument
def CheckChangeOnUpload(input_api, output_api):
  return []

def CheckChangeOnCommit(input_api, output_api):
  return CommonChecks(input_api, output_api)
<|endoftext|>


<commit_before># Copyright 2008 (C) Nicira, Inc.
#
# This file is part of NOX.
#
# NOX is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# NOX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with NOX.  If not, see <http://www.gnu.org/licenses/>.
#======================================================================
#
#                          IPv4 Header Format
#
#    0                   1                   2                   3
#    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |Version|  IHL  |Type of Service|          Total Length         |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |         Identification        |Flags|      Fragment Offset    |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |  Time to Live |    Protocol   |         Header Checksum       |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                       Source Address                          |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                    Destination Address                        |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                    Options                    |    Padding    |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#
#======================================================================

import struct
import time
from packet_utils       import *
from tcp import *
from udp import *
from icmp import *

from packet_base import packet_base

from pox.lib.addresses import IPAddr

IP_ANY = IPAddr("0.0.0.0")
IP_BROADCAST = IPAddr("255.255.25.255")

class ipv4(packet_base):
    "IP packet struct"

    MIN_LEN = 20

    IPv4 = 4
    ICMP_PROTOCOL = 1
    TCP_PROTOCOL  = 6
    UDP_PROTOCOL  = 17

    DF_FLAG = 0x02
    MF_FLAG = 0x01

    ip_id = 0

    def __init__(self, arr=None, prev=None):

        self.prev = prev

        self.v     = 4
        self.hl    = ipv4.MIN_LEN / 4
        self.tos   = 0
        self.iplen = ipv4.MIN_LEN
        self.id    = int(time.time())
        ipv4.ip_id += 1
        self.flags = 0
        self.frag  = 0
        self.ttl   = 64
        self.protocol = 0
        self.csum  = 0
        self.srcip = IP_ANY
        self.dstip = IP_ANY
        self.next  = ''

        if arr != None:
            assert(type(arr) == bytes)
            self.arr = arr
            self.parse()

    def __str__(self):
        if self.parsed == False:
            return ""

        s = ''.join(('(','[v:'+str(self.v),'hl:'+str(self.hl),\
                         'l:', str(self.iplen),'t:', \
                         str(self.ttl), ']', ipproto_to_str(self.protocol), \
                         ' cs:', '%x' %self.csum, '[',str(self.srcip), \
                         '>', str(self.dstip),'])'))
        if self.next == None:
            return s
        return ''.join((s, str(self.next)))

    def parse(self):
        dlen = len(self.arr)
        if dlen < ipv4.MIN_LEN:
            self.msg('warning IP packet data too short to parse header: data len %u' % dlen)
            return

        (vhl, self.tos, self.iplen, self.id, self.frag, self.ttl,
            self.protocol, self.csum, self.srcip, self.dstip) \
             = struct.unpack('!BBHHHBBHII', self.arr[:ipv4.MIN_LEN])

        self.v = vhl >> 4
        self.hl = vhl & 0x0f

        self.flags = self.frag >> 13
        self.frag  = self.frag & 0x1f

        if self.v != ipv4.IPv4:
            self.msg('ip parse) warning IP version %u not IPv4' % self.v)
            return
        elif self.hl < 5:
            self.msg('(ip parse) warning IP header %u longer than len %u' \
                        % (self.hl, self.iplen))
            return
        elif self.iplen < ipv4.MIN_LEN:
            self.msg('(ip parse) warning invalid IP len %u' % self.iplen)
            return
        elif (self.hl * 4) >= self.iplen or (self.hl * 4) > dlen:
            self.msg('(ip parse) warning IP header %u longer than len %u' \
                        % (self.hl, self.iplen))
            return

        self.dstip = IPAddr(self.dstip)
        self.srcip = IPAddr(self.srcip)

        # At this point, we are reasonably certain that we have an IP
        # packet
        self.parsed = True

        length = self.iplen
        if length > dlen:
            length = dlen # Clamp to what we've got
        if self.protocol == ipv4.UDP_PROTOCOL:
            self.next = udp(arr=self.arr[self.hl*4:length], prev=self)
        elif self.protocol == ipv4.TCP_PROTOCOL:
            self.next = tcp(arr=self.arr[self.hl*4:length], prev=self)
        elif self.protocol == ipv4.ICMP_PROTOCOL:
            self.next = icmp(arr=self.arr[self.hl*4:length], prev=self)
        elif dlen < self.iplen:
            self.msg('(ip parse) warning IP packet data shorter than IP len: %u < %u' % (dlen, self.iplen))
        else:
            self.next =  self.arr[self.hl*4:length]

        if isinstance(self.next, packet_base) and not self.next.parsed:
            self.next =  self.arr[self.hl*4:length]

    def checksum(self):
        data = struct.pack('!BBHHHBBHII', (self.v << 4) + self.hl, self.tos, \
                                 self.iplen, self.id, \
                                 (self.flags << 13) | self.frag, self.ttl, \
                                 self.protocol, 0, self.srcip.toUnsigned(), \
                                 self.dstip.toUnsigned())
        return checksum(data, 0)

    def hdr(self):
        self.csum = self.checksum()
        return struct.pack('!BBHHHBBHII', (self.v << 4) + self.hl, self.tos, \
                                 self.iplen, self.id, (self.flags << 13) | self.frag, self.ttl, \
                                 self.protocol, self.csum, self.srcip.toUnsigned(), \
                                 self.dstip.toUnsigned())

<commit_msg>packet: Fix IP identification field<commit_after># Copyright 2008 (C) Nicira, Inc.
#
# This file is part of NOX.
#
# NOX is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# NOX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with NOX.  If not, see <http://www.gnu.org/licenses/>.
#======================================================================
#
#                          IPv4 Header Format
#
#    0                   1                   2                   3
#    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |Version|  IHL  |Type of Service|          Total Length         |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |         Identification        |Flags|      Fragment Offset    |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |  Time to Live |    Protocol   |         Header Checksum       |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                       Source Address                          |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                    Destination Address                        |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#   |                    Options                    |    Padding    |
#   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
#
#======================================================================

import struct
import time
from packet_utils       import *
from tcp import *
from udp import *
from icmp import *

from packet_base import packet_base

from pox.lib.addresses import IPAddr

IP_ANY = IPAddr("0.0.0.0")
IP_BROADCAST = IPAddr("255.255.25.255")

class ipv4(packet_base):
    "IP packet struct"

    MIN_LEN = 20

    IPv4 = 4
    ICMP_PROTOCOL = 1
    TCP_PROTOCOL  = 6
    UDP_PROTOCOL  = 17

    DF_FLAG = 0x02
    MF_FLAG = 0x01

    ip_id = int(time.time())

    def __init__(self, arr=None, prev=None):

        self.prev = prev

        self.v     = 4
        self.hl    = ipv4.MIN_LEN / 4
        self.tos   = 0
        self.iplen = ipv4.MIN_LEN
        ipv4.ip_id = (ipv4.ip_id + 1) & 0xffff
        self.id    = ipv4.ip_id
        self.flags = 0
        self.frag  = 0
        self.ttl   = 64
        self.protocol = 0
        self.csum  = 0
        self.srcip = IP_ANY
        self.dstip = IP_ANY
        self.next  = ''

        if arr != None:
            assert(type(arr) == bytes)
            self.arr = arr
            self.parse()

    def __str__(self):
        if self.parsed == False:
            return ""

        s = ''.join(('(','[v:'+str(self.v),'hl:'+str(self.hl),\
                         'l:', str(self.iplen),'t:', \
                         str(self.ttl), ']', ipproto_to_str(self.protocol), \
                         ' cs:', '%x' %self.csum, '[',str(self.srcip), \
                         '>', str(self.dstip),'])'))
        if self.next == None:
            return s
        return ''.join((s, str(self.next)))

    def parse(self):
        dlen = len(self.arr)
        if dlen < ipv4.MIN_LEN:
            self.msg('warning IP packet data too short to parse header: data len %u' % dlen)
            return

        (vhl, self.tos, self.iplen, self.id, self.frag, self.ttl,
            self.protocol, self.csum, self.srcip, self.dstip) \
             = struct.unpack('!BBHHHBBHII', self.arr[:ipv4.MIN_LEN])

        self.v = vhl >> 4
        self.hl = vhl & 0x0f

        self.flags = self.frag >> 13
        self.frag  = self.frag & 0x1f

        if self.v != ipv4.IPv4:
            self.msg('ip parse) warning IP version %u not IPv4' % self.v)
            return
        elif self.hl < 5:
            self.msg('(ip parse) warning IP header %u longer than len %u' \
                        % (self.hl, self.iplen))
            return
        elif self.iplen < ipv4.MIN_LEN:
            self.msg('(ip parse) warning invalid IP len %u' % self.iplen)
            return
        elif (self.hl * 4) >= self.iplen or (self.hl * 4) > dlen:
            self.msg('(ip parse) warning IP header %u longer than len %u' \
                        % (self.hl, self.iplen))
            return

        self.dstip = IPAddr(self.dstip)
        self.srcip = IPAddr(self.srcip)

        # At this point, we are reasonably certain that we have an IP
        # packet
        self.parsed = True

        length = self.iplen
        if length > dlen:
            length = dlen # Clamp to what we've got
        if self.protocol == ipv4.UDP_PROTOCOL:
            self.next = udp(arr=self.arr[self.hl*4:length], prev=self)
        elif self.protocol == ipv4.TCP_PROTOCOL:
            self.next = tcp(arr=self.arr[self.hl*4:length], prev=self)
        elif self.protocol == ipv4.ICMP_PROTOCOL:
            self.next = icmp(arr=self.arr[self.hl*4:length], prev=self)
        elif dlen < self.iplen:
            self.msg('(ip parse) warning IP packet data shorter than IP len: %u < %u' % (dlen, self.iplen))
        else:
            self.next =  self.arr[self.hl*4:length]

        if isinstance(self.next, packet_base) and not self.next.parsed:
            self.next =  self.arr[self.hl*4:length]

    def checksum(self):
        data = struct.pack('!BBHHHBBHII', (self.v << 4) + self.hl, self.tos, \
                                 self.iplen, self.id, \
                                 (self.flags << 13) | self.frag, self.ttl, \
                                 self.protocol, 0, self.srcip.toUnsigned(), \
                                 self.dstip.toUnsigned())
        return checksum(data, 0)

    def hdr(self):
        self.csum = self.checksum()
        return struct.pack('!BBHHHBBHII', (self.v << 4) + self.hl, self.tos, \
                                 self.iplen, self.id, (self.flags << 13) | self.frag, self.ttl, \
                                 self.protocol, self.csum, self.srcip.toUnsigned(), \
                                 self.dstip.toUnsigned())

<|endoftext|>


<commit_before>"""
==========================================================
Morph source estimates from one subject to another subject
==========================================================

A source estimate from a given subject 'sample' is morphed
to the anatomy of another subject 'fsaverage'. The output
is a source estimate defined on the anatomy of 'fsaverage'

"""
# Authors: Alexandre Gramfort <gramfort@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#
# License: BSD (3-clause)

print __doc__

import mne
from mne.datasets import sample

data_path = sample.data_path('..')

subject_from = 'sample'
subject_to = 'fsaverage'

fname = data_path + '/MEG/sample/sample_audvis-meg'
src_fname = data_path + '/MEG/sample/sample_audvis-meg-oct-6-fwd.fif'

# Read input stc file
stc_from = mne.read_source_estimate(fname)
# Morph using one method (supplying the vertices in fsaverage's source
# space makes it faster). Note that for any generic subject, you could do:
#     vertices_to = mne.grade_to_vertices(subject_to, grade=5)
# But fsaverage's source space was set up so we can just do this:
vertices_to = [np.arange(10242), np.arange(10242)]
stc_to = mne.morph_data(subject_from, subject_to, stc_from, n_jobs=1,
                        grade=vertices_to)
stc_to.save('%s_audvis-meg' % subject_to)

# Morph using another method -- useful if you're going to do a lot of the
# same inter-subject morphing operations; you could save and load morph_mat
morph_mat = mne.compute_morph_matrix(subject_from, subject_to, stc_from.vertno,
                                     vertices_to)
stc_to_2 = mne.morph_data_precomputed(subject_from, subject_to,
                                      stc_from, vertices_to, morph_mat)
stc_to_2.save('%s_audvis-meg_2' % subject_to)

# View source activations
import pylab as pl
pl.plot(stc_from.times, stc_from.data.mean(axis=0), 'r', label='from')
pl.plot(stc_to.times, stc_to.data.mean(axis=0), 'b', label='to')
pl.plot(stc_to_2.times, stc_to.data.mean(axis=0), 'g', label='to_2')
pl.xlabel('time (ms)')
pl.ylabel('Mean Source amplitude')
pl.legend()
pl.show()<commit_msg>FIX: Import numpy<commit_after>"""
==========================================================
Morph source estimates from one subject to another subject
==========================================================

A source estimate from a given subject 'sample' is morphed
to the anatomy of another subject 'fsaverage'. The output
is a source estimate defined on the anatomy of 'fsaverage'

"""
# Authors: Alexandre Gramfort <gramfort@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#
# License: BSD (3-clause)

print __doc__

import mne
import numpy as np
from mne.datasets import sample

data_path = sample.data_path('..')

subject_from = 'sample'
subject_to = 'fsaverage'

fname = data_path + '/MEG/sample/sample_audvis-meg'
src_fname = data_path + '/MEG/sample/sample_audvis-meg-oct-6-fwd.fif'

# Read input stc file
stc_from = mne.read_source_estimate(fname)
# Morph using one method (supplying the vertices in fsaverage's source
# space makes it faster). Note that for any generic subject, you could do:
#     vertices_to = mne.grade_to_vertices(subject_to, grade=5)
# But fsaverage's source space was set up so we can just do this:
vertices_to = [np.arange(10242), np.arange(10242)]
stc_to = mne.morph_data(subject_from, subject_to, stc_from, n_jobs=1,
                        grade=vertices_to)
stc_to.save('%s_audvis-meg' % subject_to)

# Morph using another method -- useful if you're going to do a lot of the
# same inter-subject morphing operations; you could save and load morph_mat
morph_mat = mne.compute_morph_matrix(subject_from, subject_to, stc_from.vertno,
                                     vertices_to)
stc_to_2 = mne.morph_data_precomputed(subject_from, subject_to,
                                      stc_from, vertices_to, morph_mat)
stc_to_2.save('%s_audvis-meg_2' % subject_to)

# View source activations
import pylab as pl
pl.plot(stc_from.times, stc_from.data.mean(axis=0), 'r', label='from')
pl.plot(stc_to.times, stc_to.data.mean(axis=0), 'b', label='to')
pl.plot(stc_to_2.times, stc_to.data.mean(axis=0), 'g', label='to_2')
pl.xlabel('time (ms)')
pl.ylabel('Mean Source amplitude')
pl.legend()
pl.show()<|endoftext|>


<commit_before># Copyright 2011,2012,2013 James McCauley
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
A utility module for handling some mundane parts of ARP
"""

"""
TODO
----
arp_responder should be refactored to use this.  Also, it should be possible
to have a simple ARP learner which keeps an ARP table without responding...
"""

from pox.core import core
import pox
log = core.getLogger()

from pox.lib.packet.ethernet import ethernet, ETHER_BROADCAST
from pox.lib.packet.arp import arp
from pox.lib.addresses import EthAddr, IPAddr
from pox.lib.util import dpid_to_str, str_to_bool
from pox.lib.revent import EventHalt, Event, EventMixin

import pox.openflow.libopenflow_01 as of

def send_arp_reply (reply_to, mac, src_mac = None, src_ip = None):
  # reply_to should be a PacketIn event
  arpp = reply_to.parsed.find('arp')
  mac = EthAddr(mac)
  if src_mac is None:
    src_mac = mac
  else:
    src_mac = EthAddr(src_mac)
  r = arp()
  r.opcode = r.REPLY
  r.hwdst = arpp.hwsrc
  r.protodst = arpp.protosrc
  r.hwsrc = EthAddr(src_mac)
  r.protosrc = IPAddr("0.0.0.0") if src_ip is None else IPAddr(src_ip)
  e = ethernet(type=ethernet.ARP_TYPE, src=src_mac, dst=r.hwdst)
  e.payload = r
  msg = of.ofp_packet_out()
  msg.data = e.pack()
  msg.actions.append(of.ofp_action_output(port = reply_to.port))
  msg.in_port = of.OFPP_NONE
  reply_to.connection.send(msg)

def send_arp_request (connection, ip, port = of.OFPP_FLOOD,
                      src_mac = None, src_ip = None):
  if src_mac is None:
    src_mac = _dpid_to_mac(connection.dpid)
  else:
    src_mac = EthAddr(src_mac)
  r = arp()
  r.opcode = r.REQUEST
  r.hwdst = ETHER_BROADCAST
  r.protodst = IPAddr(ip)
  r.hwsrc = src_mac
  r.protosrc = IPAddr("0.0.0.0") if src_ip is None else IPAddr(src_ip)
  e = ethernet(type=ethernet.ARP_TYPE, src=src_mac, dst=r.hwdst)
  e.payload = r
  msg = of.ofp_packet_out()
  msg.data = e.pack()
  msg.actions.append(of.ofp_action_output(port = port))
  msg.in_port = of.OFPP_NONE
  connection.send(msg)

def _dpid_to_mac (dpid):
  # Should maybe look at internal port MAC instead?
  return EthAddr("%012x" % (dpid & 0xffFFffFFffFF,))

class ARPRequest (Event):
  @property
  def dpid (self):
    return self.connection.dpid

  def __str__ (self):
    return "ARPRequest for %s on %s"  % (self.ip, dpid_to_str(self.dpid))

  def __init__ (self, con, arpp, reply_from, eat_packet, port):
    super(ARPRequest,self).__init__()
    self.connection = con
    self.request = arpp # ARP packet
    self.reply_from = reply_from # MAC
    self.eat_packet = eat_packet
    self.port = port

    self.ip = arpp.protosrc
    self.reply = None # Set to desired EthAddr

class ARPReply (Event):
  @property
  def dpid (self):
    return self.connection.dpid

  def __str__ (self):
    return "ARPReply for %s on %s"  % (self.reply.protodst,
                                       dpid_to_str(self.dpid))

  def __init__ (self, con, arpp, eat_packet, port):
    super(ARPReply,self).__init__()
    self.connection = con
    self.reply = arpp
    self.eat_packet = eat_packet
    self.port = port

class ARPHelper (EventMixin):
  _eventMixin_events = set([ARPRequest,ARPReply])
  _rule_priority = 0x7000 # Pretty high

  def __init__ (self, no_flow, eat_packets):
    core.addListeners(self)
    self._install_flow = not no_flow
    self.eat_packets = eat_packets

  def send_arp_request (self, *args, **kw):
    return send_arp_request(*args, **kw)

  def send_arp_reply (self, *args, **kw):
    return send_arp_reply(*args, **kw)

  def _handle_GoingUpEvent (self, event):
    core.openflow.addListeners(self)
    log.debug("Up...")

  def _handle_ConnectionUp (self, event):
    if self._install_flow:
      fm = of.ofp_flow_mod()
      fm.priority = self._rule_priority
      fm.match.dl_type = ethernet.ARP_TYPE
      fm.actions.append(of.ofp_action_output(port=of.OFPP_CONTROLLER))
      event.connection.send(fm)

  def _handle_PacketIn (self, event):
    dpid = event.connection.dpid
    inport = event.port
    packet = event.parsed

    a = packet.find('arp')
    if not a: return

    if a.prototype != arp.PROTO_TYPE_IP:
      return

    if a.hwtype != arp.HW_TYPE_ETHERNET:
      return

    if a.opcode == arp.REQUEST:
      log.debug("%s ARP request %s => %s", dpid_to_str(dpid),
                a.protosrc, a.protodst)

      ev = ARPRequest(event.connection,a,_dpid_to_mac(dpid),self.eat_packets,
          inport)
      self.raiseEvent(ev)
      if ev.reply is not None:
        r = arp()
        r.hwtype = a.hwtype
        r.prototype = a.prototype
        r.hwlen = a.hwlen
        r.protolen = a.protolen
        r.opcode = arp.REPLY
        r.hwdst = a.hwsrc
        r.protodst = a.protosrc
        r.protosrc = a.protodst
        r.hwsrc = EthAddr(ev.reply)
        e = ethernet(type=packet.type, src=ev.reply_from, dst=a.hwsrc)
        e.payload = r
        log.debug("%s answering ARP for %s" % (dpid_to_str(dpid),
            str(r.protosrc)))
        msg = of.ofp_packet_out()
        msg.data = e.pack()
        msg.actions.append(of.ofp_action_output(port =
                                                of.OFPP_IN_PORT))
        msg.in_port = inport
        event.connection.send(msg)
        return EventHalt if ev.eat_packet else None

    elif a.opcode == arp.REPLY:
      log.debug("%s ARP reply %s => %s", dpid_to_str(dpid),
                a.protosrc, a.hwsrc)

      ev = ARPReply(event.connection,a,self.eat_packets,inport)
      self.raiseEvent(ev)
      return EventHalt if ev.eat_packet else None

    return EventHalt if self.eat_packets else None

def launch (no_flow=False, eat_packets=True):
  core.registerNew(ARPHelper, str_to_bool(no_flow), str_to_bool(eat_packets))
<commit_msg>misc.arp_helper: More options for source MAC<commit_after># Copyright 2011,2012,2013 James McCauley
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
A utility module for handling some mundane parts of ARP
"""

"""
TODO
----
arp_responder should be refactored to use this.  Also, it should be possible
to have a simple ARP learner which keeps an ARP table without responding...
"""

from pox.core import core
import pox
log = core.getLogger()

from pox.lib.packet.ethernet import ethernet, ETHER_BROADCAST
from pox.lib.packet.arp import arp
from pox.lib.addresses import EthAddr, IPAddr
from pox.lib.util import dpid_to_str, str_to_bool
from pox.lib.revent import EventHalt, Event, EventMixin

import pox.openflow.libopenflow_01 as of

def send_arp_reply (reply_to, mac, src_mac = None, src_ip = None):
  """
  Send an ARP reply.

  src_mac can be None to use the "DPID MAC" or True to use the port Mac.
    (or it can be an EthAddr)
  """
  # reply_to should be a PacketIn event
  arpp = reply_to.parsed.find('arp')
  mac = EthAddr(mac)
  if src_mac is None:
    #src_mac = mac # Used to be this ???
    src_mac = reply_to.connection.eth_addr
  elif src_mac is True:
    src_mac = reply_to.connection.ports[reply_to.port].hw_addr
  else:
    src_mac = EthAddr(src_mac)
  r = arp()
  r.opcode = r.REPLY
  r.hwdst = arpp.hwsrc
  r.protodst = arpp.protosrc
  r.hwsrc = EthAddr(src_mac)
  r.protosrc = IPAddr("0.0.0.0") if src_ip is None else IPAddr(src_ip)
  e = ethernet(type=ethernet.ARP_TYPE, src=src_mac, dst=r.hwdst)
  e.payload = r
  msg = of.ofp_packet_out()
  msg.data = e.pack()
  msg.actions.append(of.ofp_action_output(port = reply_to.port))
  msg.in_port = of.OFPP_NONE
  reply_to.connection.send(msg)

def send_arp_request (connection, ip, port = of.OFPP_FLOOD,
                      src_mac = None, src_ip = None):
  """
  Send an ARP request

  src_mac can be None to use the "DPID MAC" or True to use the port Mac.
    (or it can be an EthAddr)
  """
  if src_mac is None:
    src_mac = connection.eth_addr
  elif src_mac is True:
    if port in (of.OFPP_FLOOD, of.OFPP_ALL):
      for p in connection.ports.values():
        if p.config & OFPPC_NO_FLOOD:
          if port == of.ofPP_FLOOD:
            continue
        if p.port_no < 0: continue
        if p.port_no > of.OFPP_MAX: continue # Off by one?
        send_arp_request(connection, ip, p.port_no,
                         src_mac=p.hw_addr, src_ip=src_ip)
      return
    src_mac = connection.ports[port].hw_addr
  else:
    src_mac = EthAddr(src_mac)
  r = arp()
  r.opcode = r.REQUEST
  r.hwdst = ETHER_BROADCAST
  r.protodst = IPAddr(ip)
  r.hwsrc = src_mac
  r.protosrc = IPAddr("0.0.0.0") if src_ip is None else IPAddr(src_ip)
  e = ethernet(type=ethernet.ARP_TYPE, src=src_mac, dst=r.hwdst)
  e.payload = r
  msg = of.ofp_packet_out()
  msg.data = e.pack()
  msg.actions.append(of.ofp_action_output(port = port))
  msg.in_port = of.OFPP_NONE
  connection.send(msg)

class ARPRequest (Event):
  @property
  def dpid (self):
    return self.connection.dpid

  def __str__ (self):
    return "ARPRequest for %s on %s"  % (self.ip, dpid_to_str(self.dpid))

  def __init__ (self, con, arpp, reply_from, eat_packet, port):
    super(ARPRequest,self).__init__()
    self.connection = con
    self.request = arpp # ARP packet
    self.reply_from = reply_from # MAC
    self.eat_packet = eat_packet
    self.port = port

    self.ip = arpp.protosrc
    self.reply = None # Set to desired EthAddr

class ARPReply (Event):
  @property
  def dpid (self):
    return self.connection.dpid

  def __str__ (self):
    return "ARPReply for %s on %s"  % (self.reply.protodst,
                                       dpid_to_str(self.dpid))

  def __init__ (self, con, arpp, eat_packet, port):
    super(ARPReply,self).__init__()
    self.connection = con
    self.reply = arpp
    self.eat_packet = eat_packet
    self.port = port

_default_src_mac = object()

class ARPHelper (EventMixin):
  _eventMixin_events = set([ARPRequest,ARPReply])
  _rule_priority = 0x7000 # Pretty high

  def __init__ (self, no_flow, eat_packets, use_port_mac = False):
    core.addListeners(self)
    self._install_flow = not no_flow
    self.eat_packets = eat_packets
    self.use_port_mac = use_port_mac

  def send_arp_request (self, connection, ip, port = of.OFPP_FLOOD,
                        src_mac = _default_src_mac, src_ip = None):
    if src_mac is _default_src_mac:
      src_mac = True if self.use_port_mac else None
    return send_arp_request(connection, ip, port, src_mac, src_ip)

  def send_arp_reply (self, reply_to, mac,
                      src_mac = _default_src_mac, src_ip = None):
    if src_mac is _default_src_mac:
      src_mac = True if self.use_port_mac else None
    return send_arp_reply(reply_to, mac, src_mac, src_ip)

  def _handle_GoingUpEvent (self, event):
    core.openflow.addListeners(self)
    log.debug("Up...")

  def _handle_ConnectionUp (self, event):
    if self._install_flow:
      fm = of.ofp_flow_mod()
      fm.priority = self._rule_priority
      fm.match.dl_type = ethernet.ARP_TYPE
      fm.actions.append(of.ofp_action_output(port=of.OFPP_CONTROLLER))
      event.connection.send(fm)

  def _handle_PacketIn (self, event):
    dpid = event.connection.dpid
    inport = event.port
    packet = event.parsed

    a = packet.find('arp')
    if not a: return

    if a.prototype != arp.PROTO_TYPE_IP:
      return

    if a.hwtype != arp.HW_TYPE_ETHERNET:
      return

    if a.opcode == arp.REQUEST:
      log.debug("%s ARP request %s => %s", dpid_to_str(dpid),
                a.protosrc, a.protodst)

      if self.use_port_mac:
        src_mac = event.connection.ports[inport].hw_addr
      else:
        src_mac = event.connection.eth_addr
      ev = ARPRequest(event.connection, a, src_mac,
                      self.eat_packets, inport)
      self.raiseEvent(ev)
      if ev.reply is not None:
        r = arp()
        r.hwtype = a.hwtype
        r.prototype = a.prototype
        r.hwlen = a.hwlen
        r.protolen = a.protolen
        r.opcode = arp.REPLY
        r.hwdst = a.hwsrc
        r.protodst = a.protosrc
        r.protosrc = a.protodst
        r.hwsrc = EthAddr(ev.reply)
        e = ethernet(type=packet.type, src=ev.reply_from, dst=a.hwsrc)
        e.payload = r
        log.debug("%s answering ARP for %s" % (dpid_to_str(dpid),
            str(r.protosrc)))
        msg = of.ofp_packet_out()
        msg.data = e.pack()
        msg.actions.append(of.ofp_action_output(port =
                                                of.OFPP_IN_PORT))
        msg.in_port = inport
        event.connection.send(msg)
        return EventHalt if ev.eat_packet else None

    elif a.opcode == arp.REPLY:
      log.debug("%s ARP reply %s => %s", dpid_to_str(dpid),
                a.protosrc, a.hwsrc)

      ev = ARPReply(event.connection,a,self.eat_packets,inport)
      self.raiseEvent(ev)
      return EventHalt if ev.eat_packet else None

    return EventHalt if self.eat_packets else None

def launch (no_flow=False, eat_packets=True, use_port_mac = False):
  core.registerNew(ARPHelper, str_to_bool(no_flow), str_to_bool(eat_packets),
                   str_to_bool(use_port_mac))
<|endoftext|>


<commit_before>""" Text images of expressions. """

import sys

from . import constants as c
from . import declarations
from . import tags as t

# Helpers
# ============================================================================

def perror(message):
    """ Shorthand to print to `stderr`. """
    print(message, file=sys.stderr)

def error(message):
    """ `perror` and `sys.exit(1)`. """
    perror(message)
    sys.exit(1)

# Sort image
# ============================================================================

def sort_image(node, paren_if_fun=False):
    """ Image of s2xxx_srt and S2RTfun[0][n]. """
    # Node is a {S2RTbas}|{S2RTfun}
    if t.S2RTBAS in node:
        result = node[t.S2RTBAS][0]
    elif t.S2RTFUN in node:
        node = node[t.S2RTFUN]
        inputs = node[0]
        output = node[1]
        result = ""
        if paren_if_fun:
            result += "("
        if inputs:
            marny_args = len(inputs) > 1
            if marny_args:
                result += "("
            first = True
            for item in inputs:
                if not first:
                    result += ", "
                result += sort_image(item, paren_if_fun=not marny_args)
                first = False
            if marny_args:
                result += ")"
            elif first:
                result += "()"
            result += " -> "
        result += sort_image(output)
        if paren_if_fun:
            result += ")"
    elif t.S2RTTUP in node:
        node = node[t.S2RTTUP]
        items = node[0]
        result = "("
        first = True
        for item in items:
            if not first:
                result += ", "
            result += sort_image(item)
            first = False
        result += ")"
    else:
        error("Unknown sort expression")
    return result

# Quantified expression image
# ============================================================================

def quantified_exp_image(node, key_image, open_close):
    """ Image of an S2Eexi or of an S2Euni, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    (opn, close) = open_close  # Two paired characters.
    variables = node[0]
    predicats = node[1]
    expression = node[2]
    result = opn
    first = True
    for variable in variables:
        if not first:
            result += "; "
        result += s2var_image(variable[t.S2VAR_STAMP])
        first = False
    if predicats:
        if variables:
            result += " | "
        first = True
        for predicat in predicats:
            if not first:
                result += "; "
            result += image(predicat[key])
            first = False
    result += close
    result += " "
    result += image(expression[key])
    return result

# Type image
# ============================================================================

# Special cases
# ----------------------------------------------------------------------------

def s2var_image(stamp):
    """ Static variable image. """
    den = declarations.get_def(t.S2VAR_STAMP, stamp)
    if den is None:
        return "*ERROR*1*"
    return den.name + ": " + sort_image(den.sort)

def s2ecst_image(node, for_type):
    """ Image of a S2Ecst, either as type or sort.

    Type or sort, depending on `for_type`.

    """
    stamp = node[0][t.S2CST_STAMP]
    den = declarations.get_def(t.S2CST_STAMP, stamp)
    if den is None:
        return "*ERROR*2*"
    if for_type:
        return den.name
    return sort_image(den.sort)

def s2evar_image(node, for_type):
    """ Image of a S2Evar, either as type or sort.

    Type or sort, depending on `for_type`.

    """
    stamp = node[0][t.S2VAR_STAMP]
    den = declarations.get_def(t.S2VAR_STAMP, stamp)
    if den is None:
        return "*ERROR*3*"
    if for_type:
        return den.name
    return sort_image(den.sort)

# Specific
# ----------------------------------------------------------------------------

def s2eapp_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Eapp, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    function = node[0]
    arguments = node[1]
    result = ""
    if paren_if_app:
        result += "("
    result += image(function[key], paren_if_fun=True)
    result += "("
    first = True
    for item in arguments:
        if not first:
            result += ", "
        result += image(item[key])
        first = False
    result += ")"
    if paren_if_app:
        result += ")"
    return result

def s2eexi_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Eexi, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    return quantified_exp_image(node, key_image, open_close=("[", "]"))

def s2efun_image(node, key_image, paren_if_fun, _paren_if_app):
    """ Image of a S2Efun, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    inputs = node[1]
    output = node[2]
    result = ""
    if paren_if_fun:
        result += "("
    marny_args = len(inputs) > 1
    if marny_args:
        result += "("
    first = True
    for item in inputs:
        if not first:
            result += ", "
        result += image(item[key], not marny_args)
        first = False
    if marny_args:
        result += ")"
    elif first:
        result += "()"
    result += " -> "
    result += image(output[key])
    if paren_if_fun:
        result += ")"
    return result

def s2erefarg_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of a S2Erefarg, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    passing_style = node[0]
    if passing_style == c.BY_VALUE:
        prefix = "!"
    elif passing_style == c.BY_REFERENCE:
        prefix = "&"
    else:
        error("Unknown argument passing style: %i" % passing_style)
    return prefix + image(node[1][key], paren_if_fun=True, paren_if_app=True)

def s2etop_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etop, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    view_status = node[0]
    if view_status == c.UNINITIALIZED:
        postfix = "?"
    elif view_status == c.INITIALIZED:
        postfix = "?!"
    else:
        error("Unknown view status: %i" % view_status)
    result = ""
    if paren_if_app:
        result += "("
    result += image(node[1][key], paren_if_app=True)
    result += postfix
    if paren_if_app:
        result += ")"
    return result

def s2etyarr_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etyarr, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    element_type = node[0]
    dimmensions = node[1]
    result = ""
    if paren_if_app:
        result += "("
    result += "@["
    result += image(element_type[key])
    result += "]"
    result += "["
    first = True
    for dimmension in dimmensions:
        if not first:
            result += ", "
        result += image(dimmension[key])
        first = False
    result += "]"
    if paren_if_app:
        result += ")"
    return result

def s2etyrec_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etyrec, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    kind_tag = node[0]
    flat = True
    if t.TYRECKINDBOX in kind_tag or t.TYRECKINDBOX_LIN in kind_tag:
        flat = False
    labels = node[2]
    result = ""
    if paren_if_app:
        result += "("
    if flat:
        result += "@{"
    else:
        result += "'{"
    first = True
    for label in labels:
        if not first:
            result += ", "
        label = label[t.SL0ABELED]
        label_name = label[0]
        if t.LABINT in label_name:
            result += str(label_name[t.LABINT])
        elif t.LABSYM in label_name:
            result += label_name[t.LABSYM]
        else:
            error("Unknown label name type")
        result += "="
        result += image(label[2][key])
        first = False
    result += "}"
    if paren_if_app:
        result += ")"
    return result

def s2euni_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Euni, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    return quantified_exp_image(node, key_image, open_close=("{", "}"))

def s2ewthtype_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Ewthtype, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    return image(node[0][key])

# Dispatch table
# ----------------------------------------------------------------------------

# There five special cases not in this table.

DISPATCH = {
    t.S2EAPP: s2eapp_image,
    t.S2EEXI: s2eexi_image,
    t.S2EFUN: s2efun_image,
    t.S2EREFARG: s2erefarg_image,
    t.S2ETOP: s2etop_image,
    t.S2ETYARR: s2etyarr_image,
    t.S2ETYREC: s2etyrec_image,
    t.S2EUNI: s2euni_image,
    t.S2EWTHTYPE: s2ewthtype_image}

# Main
# ============================================================================

def s2e_image(node, for_type, paren_if_fun=False, paren_if_app=False):
    """ Image of s2exp_node, either as type or sort. """
    if for_type:
        key_image = (t.S2EXP_NODE, type_image)
    else:
        key_image = (t.S2EXP_SRT, sort_image)

    keys = list(node.keys())
    assert len(keys) == 1
    key = keys[0]
    sub_node = node[key]

    if key == t.S2ECST:
        result = s2ecst_image(sub_node, for_type)
    elif key == t.S2EVAR:
        result = s2evar_image(sub_node, for_type)
    elif key == t.S2EEXTKIND:
        result = sub_node[0] if for_type else "?"
    elif key == t.S2EEXTYPE:
        result = sub_node[0] if for_type else "?"
    elif key == t.S2EINTINF:
        result = sub_node[0] if for_type else "?"
    elif key in DISPATCH:
        method = DISPATCH[key]
        result = method(sub_node, key_image, paren_if_fun, paren_if_app)
    else:
        result = "?"
    return result

def type_image(node, paren_if_fun=False, paren_if_app=False):
    """ Dyn image. """
    return s2e_image(node, True, paren_if_fun, paren_if_app)

def type_sorts_image(node, paren_if_fun=False, paren_if_app=False):
    """ Dyn image. """
    return s2e_image(node, False, paren_if_fun, paren_if_app)
<commit_msg>Some types are anonymous<commit_after>""" Text images of expressions. """

import sys

from . import constants as c
from . import declarations
from . import tags as t

# Helpers
# ============================================================================

def perror(message):
    """ Shorthand to print to `stderr`. """
    print(message, file=sys.stderr)

def error(message):
    """ `perror` and `sys.exit(1)`. """
    perror(message)
    sys.exit(1)

# Sort image
# ============================================================================

def sort_image(node, paren_if_fun=False):
    """ Image of s2xxx_srt and S2RTfun[0][n]. """
    # Node is a {S2RTbas}|{S2RTfun}
    if t.S2RTBAS in node:
        result = node[t.S2RTBAS][0]
    elif t.S2RTFUN in node:
        node = node[t.S2RTFUN]
        inputs = node[0]
        output = node[1]
        result = ""
        if paren_if_fun:
            result += "("
        if inputs:
            marny_args = len(inputs) > 1
            if marny_args:
                result += "("
            first = True
            for item in inputs:
                if not first:
                    result += ", "
                result += sort_image(item, paren_if_fun=not marny_args)
                first = False
            if marny_args:
                result += ")"
            elif first:
                result += "()"
            result += " -> "
        result += sort_image(output)
        if paren_if_fun:
            result += ")"
    elif t.S2RTTUP in node:
        node = node[t.S2RTTUP]
        items = node[0]
        result = "("
        first = True
        for item in items:
            if not first:
                result += ", "
            result += sort_image(item)
            first = False
        result += ")"
    else:
        error("Unknown sort expression")
    return result

# Quantified expression image
# ============================================================================

def quantified_exp_image(node, key_image, open_close):
    """ Image of an S2Eexi or of an S2Euni, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    (opn, close) = open_close  # Two paired characters.
    variables = node[0]
    predicats = node[1]
    expression = node[2]
    result = opn
    first = True
    for variable in variables:
        if not first:
            result += "; "
        result += s2var_image(variable[t.S2VAR_STAMP])
        first = False
    if predicats:
        if variables:
            result += " | "
        first = True
        for predicat in predicats:
            if not first:
                result += "; "
            result += image(predicat[key])
            first = False
    result += close
    result += " "
    result += image(expression[key])
    return result

# Type image
# ============================================================================

# Special cases
# ----------------------------------------------------------------------------

def s2var_image(stamp):
    """ Static variable image. """
    den = declarations.get_def(t.S2VAR_STAMP, stamp)
    if den is None:
        return "*ERROR*1*"
    return den.name + ": " + sort_image(den.sort)

def s2ecst_image(node, for_type):
    """ Image of a S2Ecst, either as type or sort.

    Type or sort, depending on `for_type`.

    """
    stamp = node[0][t.S2CST_STAMP]
    den = declarations.get_def(t.S2CST_STAMP, stamp)
    if den is None:
        return "anonymous‑type"
    if for_type:
        return den.name
    return sort_image(den.sort)

def s2evar_image(node, for_type):
    """ Image of a S2Evar, either as type or sort.

    Type or sort, depending on `for_type`.

    """
    stamp = node[0][t.S2VAR_STAMP]
    den = declarations.get_def(t.S2VAR_STAMP, stamp)
    if den is None:
        return "*ERROR*3*"
    if for_type:
        return den.name
    return sort_image(den.sort)

# Specific
# ----------------------------------------------------------------------------

def s2eapp_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Eapp, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    function = node[0]
    arguments = node[1]
    result = ""
    if paren_if_app:
        result += "("
    result += image(function[key], paren_if_fun=True)
    result += "("
    first = True
    for item in arguments:
        if not first:
            result += ", "
        result += image(item[key])
        first = False
    result += ")"
    if paren_if_app:
        result += ")"
    return result

def s2eexi_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Eexi, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    return quantified_exp_image(node, key_image, open_close=("[", "]"))

def s2efun_image(node, key_image, paren_if_fun, _paren_if_app):
    """ Image of a S2Efun, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    inputs = node[1]
    output = node[2]
    result = ""
    if paren_if_fun:
        result += "("
    marny_args = len(inputs) > 1
    if marny_args:
        result += "("
    first = True
    for item in inputs:
        if not first:
            result += ", "
        result += image(item[key], not marny_args)
        first = False
    if marny_args:
        result += ")"
    elif first:
        result += "()"
    result += " -> "
    result += image(output[key])
    if paren_if_fun:
        result += ")"
    return result

def s2erefarg_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of a S2Erefarg, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    passing_style = node[0]
    if passing_style == c.BY_VALUE:
        prefix = "!"
    elif passing_style == c.BY_REFERENCE:
        prefix = "&"
    else:
        error("Unknown argument passing style: %i" % passing_style)
    return prefix + image(node[1][key], paren_if_fun=True, paren_if_app=True)

def s2etop_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etop, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    view_status = node[0]
    if view_status == c.UNINITIALIZED:
        postfix = "?"
    elif view_status == c.INITIALIZED:
        postfix = "?!"
    else:
        error("Unknown view status: %i" % view_status)
    result = ""
    if paren_if_app:
        result += "("
    result += image(node[1][key], paren_if_app=True)
    result += postfix
    if paren_if_app:
        result += ")"
    return result

def s2etyarr_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etyarr, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    element_type = node[0]
    dimmensions = node[1]
    result = ""
    if paren_if_app:
        result += "("
    result += "@["
    result += image(element_type[key])
    result += "]"
    result += "["
    first = True
    for dimmension in dimmensions:
        if not first:
            result += ", "
        result += image(dimmension[key])
        first = False
    result += "]"
    if paren_if_app:
        result += ")"
    return result

def s2etyrec_image(node, key_image, _paren_if_fun, paren_if_app):
    """ Image of an S2Etyrec, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    kind_tag = node[0]
    flat = True
    if t.TYRECKINDBOX in kind_tag or t.TYRECKINDBOX_LIN in kind_tag:
        flat = False
    labels = node[2]
    result = ""
    if paren_if_app:
        result += "("
    if flat:
        result += "@{"
    else:
        result += "'{"
    first = True
    for label in labels:
        if not first:
            result += ", "
        label = label[t.SL0ABELED]
        label_name = label[0]
        if t.LABINT in label_name:
            result += str(label_name[t.LABINT])
        elif t.LABSYM in label_name:
            result += label_name[t.LABSYM]
        else:
            error("Unknown label name type")
        result += "="
        result += image(label[2][key])
        first = False
    result += "}"
    if paren_if_app:
        result += ")"
    return result

def s2euni_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Euni, either as type or sort.

    Type or sort, depending on `key_image`.

    Of an S2Eexi or of an S2Euni, depending on `open_close` characters.

    """
    return quantified_exp_image(node, key_image, open_close=("{", "}"))

def s2ewthtype_image(node, key_image, _paren_if_fun, _paren_if_app):
    """ Image of an S2Ewthtype, either as type or sort.

    Type or sort, depending on `key_image`.

    """
    (key, image) = key_image  # `image` is a function.
    assert key == t.S2EXP_NODE or key == t.S2EXP_SRT
    return image(node[0][key])

# Dispatch table
# ----------------------------------------------------------------------------

# There five special cases not in this table.

DISPATCH = {
    t.S2EAPP: s2eapp_image,
    t.S2EEXI: s2eexi_image,
    t.S2EFUN: s2efun_image,
    t.S2EREFARG: s2erefarg_image,
    t.S2ETOP: s2etop_image,
    t.S2ETYARR: s2etyarr_image,
    t.S2ETYREC: s2etyrec_image,
    t.S2EUNI: s2euni_image,
    t.S2EWTHTYPE: s2ewthtype_image}

# Main
# ============================================================================

def s2e_image(node, for_type, paren_if_fun=False, paren_if_app=False):
    """ Image of s2exp_node, either as type or sort. """
    if for_type:
        key_image = (t.S2EXP_NODE, type_image)
    else:
        key_image = (t.S2EXP_SRT, sort_image)

    keys = list(node.keys())
    assert len(keys) == 1
    key = keys[0]
    sub_node = node[key]

    if key == t.S2ECST:
        result = s2ecst_image(sub_node, for_type)
    elif key == t.S2EVAR:
        result = s2evar_image(sub_node, for_type)
    elif key == t.S2EEXTKIND:
        result = sub_node[0] if for_type else "?"
    elif key == t.S2EEXTYPE:
        result = sub_node[0] if for_type else "?"
    elif key == t.S2EINTINF:
        result = sub_node[0] if for_type else "?"
    elif key in DISPATCH:
        method = DISPATCH[key]
        result = method(sub_node, key_image, paren_if_fun, paren_if_app)
    else:
        result = "?"
    return result

def type_image(node, paren_if_fun=False, paren_if_app=False):
    """ Dyn image. """
    return s2e_image(node, True, paren_if_fun, paren_if_app)

def type_sorts_image(node, paren_if_fun=False, paren_if_app=False):
    """ Dyn image. """
    return s2e_image(node, False, paren_if_fun, paren_if_app)
<|endoftext|>


<commit_before># encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------

"""
Calculate the Julia set for a given z <- z**2 + c with
distarray passed via command line args. Usage:
    $ python julia_distarray.py <c real component> <c imaginary component>
"""

import sys
from matplotlib import pyplot

from distarray.dist import Context, Distribution
from distarray.dist.distarray import DistArray
from distarray.dist.decorators import local

context = Context()
with context.view.sync_imports():
    import numpy

def create_complex_plane(resolution, dist, re_ax, im_ax):
    ''' Create a DistArray containing points on the complex plane.

    resolution: A 2-tuple with the number of points along Re and Im axes.
    dist: Distribution for the DistArray.
    re_ax: A 2-tuple with the (lower, upper) range of the Re axis.
    im_ax: A 2-tuple with the (lower, upper) range of the Im axis.
    '''

    @local
    def fill_complex_plane(arr, re_ax, im_ax, resolution):
        ''' Fill in points on the complex coordinate plane. '''
        # Drawing the coordinate plane directly like this is currently much
        # faster than trying to do it by indexing a distarray.
        # This may not be the most DistArray-thonic way to do this.
        re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
        im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
        for i in arr.distribution[0].global_iter:
            for j in arr.distribution[1].global_iter:
                arr.global_index[i, j] = complex(re_ax[0] + re_step*i,
                                                 im_ax[0] + im_step*j)

    # Create an empty distributed array.
    distribution = Distribution.from_shape(context,
                                           (resolution[0], resolution[1]),
                                           dist=dist)
    complex_plane = context.empty(distribution, dtype=complex)
    fill_complex_plane(complex_plane, re_ax, im_ax, resolution)
    return complex_plane

def local_julia_calc(la, c, z_max, n_max):
    ''' Calculate the number of iterations for the point to escape. '''

    @numpy.vectorize
    def julia_calc(z, c, z_max, n_max):
        ''' Use usual numpy.vectorize to apply on all the complex points. '''
        n = 0
        while abs(z) < z_max and n < n_max:
            z = z*z + c
            n += 1
        return n

    from distarray.local import LocalArray
    a = la.ndarray
    b = julia_calc(a, c, z_max=z_max, n_max=n_max)
    res = LocalArray(la.distribution, buf=b)
    rtn = proxyize(res)
    return rtn

def distributed_julia_calc(distarray, c, z_max=10, n_max=100):
    ''' Calculate the Julia set for an array of points in the complex plane. '''
    context = distarray.context
    iters_key = context.apply(local_julia_calc,
                              (distarray.key, c, z_max, n_max))
    iters_da = DistArray.from_localarrays(iters_key[0], context=context)
    return iters_da

# Grid parameters
re_ax = (-1.5, 1.5)
im_ax = (-1.5, 1.5)
dimensions = (500, 500)
# Julia set parameters, changing these is fun.
c = complex(float(sys.argv[1]), float(sys.argv[2]))
z_max = 10
n_max = 100
# Array distribution parameters
dist = {0: 'c', 1: 'c'}
dist = {0: 'b', 1: 'b'}

if __name__ == '__main__':
    # Create a distarray for the points on the complex plane.
    complex_plane = create_complex_plane(dimensions, dist, re_ax, im_ax)
    # Calculate the number of iterations to escape for each point.
    num_iters = distributed_julia_calc(complex_plane,
                                       c,
                                       z_max=z_max,
                                       n_max=n_max)
    # Draw the iteration count.
    image = num_iters.tondarray()
    pyplot.matshow(image)
    pyplot.show()
<commit_msg>Formatting improvements.<commit_after># encoding: utf-8
# ---------------------------------------------------------------------------
#  Copyright (C) 2008-2014, IPython Development Team and Enthought, Inc.
#  Distributed under the terms of the BSD License.  See COPYING.rst.
# ---------------------------------------------------------------------------

"""
Calculate the Julia set for a given z <- z**2 + c with
distarray passed via command line args. Usage:
    $ python julia_distarray.py <c real component> <c imaginary component>
"""

import sys
from matplotlib import pyplot

from distarray.dist import Context, Distribution
from distarray.dist.distarray import DistArray
from distarray.dist.decorators import local

context = Context()
with context.view.sync_imports():
    import numpy

def create_complex_plane(resolution, dist, re_ax, im_ax):
    ''' Create a DistArray containing points on the complex plane.

    resolution: A 2-tuple with the number of points along Re and Im axes.
    dist: Distribution for the DistArray.
    re_ax: A 2-tuple with the (lower, upper) range of the Re axis.
    im_ax: A 2-tuple with the (lower, upper) range of the Im axis.
    '''

    @local
    def fill_complex_plane(arr, re_ax, im_ax, resolution):
        ''' Fill in points on the complex coordinate plane. '''
        # Drawing the coordinate plane directly like this is currently much
        # faster than trying to do it by indexing a distarray.
        # This may not be the most DistArray-thonic way to do this.
        re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
        im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
        for i in arr.distribution[0].global_iter:
            for j in arr.distribution[1].global_iter:
                arr.global_index[i, j] = complex(re_ax[0] + re_step * i,
                                                 im_ax[0] + im_step * j)

    # Create an empty distributed array.
    distribution = Distribution.from_shape(context,
                                           (resolution[0], resolution[1]),
                                           dist=dist)
    complex_plane = context.empty(distribution, dtype=complex)
    fill_complex_plane(complex_plane, re_ax, im_ax, resolution)
    return complex_plane

def local_julia_calc(la, c, z_max, n_max):
    ''' Calculate the number of iterations for the point to escape.

    la: Local array of complex values whose iterations we will count.
    c: Complex number to add at each iteration.
    z_max: Magnitude of complex value that we assume goes to infinity.
    n_max: Maximum number of iterations.
    '''

    @numpy.vectorize
    def julia_calc(z, c, z_max, n_max):
        ''' Use usual numpy.vectorize to apply on all the complex points. '''
        n = 0
        while abs(z) < z_max and n < n_max:
            z = z * z + c
            n += 1
        return n

    from distarray.local import LocalArray
    a = la.ndarray
    b = julia_calc(a, c, z_max, n_max)
    res = LocalArray(la.distribution, buf=b)
    rtn = proxyize(res)
    return rtn

def distributed_julia_calc(distarray, c, z_max=10, n_max=100):
    ''' Calculate the Julia set for an array of points in the complex plane.

    distarray: DistArray of complex values whose iterations we will count.
    c: Complex number to add at each iteration.
    z_max: Magnitude of complex value that we assume goes to infinity.
    n_max: Maximum number of iterations.
    '''
    context = distarray.context
    iters_key = context.apply(local_julia_calc,
                              (distarray.key, c, z_max, n_max))
    iters_da = DistArray.from_localarrays(iters_key[0], context=context)
    return iters_da

# Grid parameters
re_ax = (-1.5, 1.5)
im_ax = (-1.5, 1.5)
dimensions = (500, 500)
# Julia set parameters, changing these is fun.
c = complex(float(sys.argv[1]), float(sys.argv[2]))
z_max = 10
n_max = 100
# Array distribution parameters
dist = {0: 'c', 1: 'c'}
dist = {0: 'b', 1: 'b'}

if __name__ == '__main__':
    # Create a distarray for the points on the complex plane.
    complex_plane = create_complex_plane(dimensions, dist, re_ax, im_ax)
    # Calculate the number of iterations to escape for each point.
    num_iters = distributed_julia_calc(complex_plane,
                                       c,
                                       z_max=z_max,
                                       n_max=n_max)
    # Draw the iteration count.
    image = num_iters.tondarray()
    pyplot.matshow(image)
    pyplot.show()
<|endoftext|>


<commit_before># Copyright 2017 James McCauley
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Reusable ARP table
"""

import time
from pox.lib.recoco import Timer
from pox.core import core
import pox.lib.packet as pkt

log = core.getLogger()

class ARPEntry (object):
  REFRESH_TIME = 60 # Refresh this often
  MIN_BACKOFF = 1
  MAX_BACKOFF = 32

  def __init__ (self, ip, mac):
    self.ts = time.time()
    self.ip = ip
    self.mac = mac
    if mac is not None:
      self.confirm()
    else:
      self.next_refresh = time.time()
      self.refresh_backoff = self.MIN_BACKOFF

  @property
  def age (self):
    return time.time() - self.ts

  def confirm (self, mac=None):
    self.next_refresh = time.time() + self.REFRESH_TIME
    self.refresh_backoff = self.MIN_BACKOFF
    if mac is not None:
      self.mac = mac

  def maybe_refresh (self):
    """
    if returns True, you should send a refresh
    """
    if self.next_refresh > time.time():
      return False
    self.next_refresh = time.time() + self.refresh_backoff
    self.refresh_backoff *= 2
    if self.refresh_backoff > self.MAX_BACKOFF:
      self.refresh_backoff = self.MAX_BACKOFF
    return True

class ARPTable (object):
  MAX_ENTRIES = 1024
  MAX_PENDING = 5

  MAX_PENDING_TIME = 2

  def __init__ (self):
    self.by_ip = {} # ip -> entry
    self.pending = [] # Packets waiting to be sent (ip,(send_args))
    self.timer = Timer(self.MAX_PENDING_TIME, self._timer_proc)

  def add_entry (self, ip, mac=None):
    """
    Add an entry

    The entry can't already exist.
    It will definitely exist after returning.
    """
    assert ip not in self.by_ip
    if len(self.by_ip) >= self.MAX_ENTRIES:
      # Sloppy, but simple.
      # Get ones with lowest age
      entries = sorted(self.by_ip.values(), key = lambda entry: entry.age)
      del entries[self.MAX_ENTRIES:]
      self.by_ip = {e.mac:e for e in entries}
    new_entry = ARPEntry(ip=ip, mac=mac)
    self.by_ip[ip] = new_entry
    return new_entry

  def send (self, eth_packet, router_ip=None, src_eth=None, src_ip=None,
            send_function=None):
    """
    Try to send a packet

    eth_packet is an ethernet object.
    src_eth is the source for any ARPs sent.
    src_ip is the source for any ARPs sent.
    If the above two are not specified, they are taken from eth_packet.
    send_function is a function which takes raw bytes to send.
    If send_function is unset, it is taken from a send_function attribute.
    """
    if send_function is None: send_function = self.send_function
    ipp = eth_packet.find("ipv4")

    if not ipp and eth_packet.type == eth_packet.IP_TYPE:
      if isinstance(eth_packet.payload, bytes):
        # Hm!  Try harder...
        ipp = pkt.ipv4(raw=eth_packet.payload)

    if not ipp or eth_packet.dst.is_multicast:
      send_function(eth_packet.pack())
      return
    if ipp.dstip == pkt.IPV4.IP_BROADCAST:
      # Would be nice to do this for any multicast...
      #ipp.dstip = router_ip # Not sure what this was about
      eth_packet.dst = pkt.ETHERNET.ETHER_BROADCAST
      send_function(eth_packet.pack())
      return

    if src_ip is None: src_ip = ipp.srcip
    if src_eth is None: src_eth = eth_packet.src

    if router_ip is not None: dstip = router_ip
    else:                     dstip = ipp.dstip

    if dstip not in self.by_ip: self.add_entry(dstip)

    e = self.by_ip[dstip]
    if e.maybe_refresh():
      # Send ARP
      self._send_arp(dstip, src_eth, src_ip,
                     send_function)

    if e.mac is not None:
      eth_packet.dst = e.mac
      send_function(eth_packet.pack())
    else:
      if len(self.pending) < self.MAX_PENDING:
        self.pending.append((dstip, (eth_packet, router_ip, src_eth, src_ip,
                                     send_function)))

  def _timer_proc (self):
    # We just blow away all the entries every interval, so on average, they
    # live for half the interval.
    del self.pending[:]

  def __del__ (self):
    if self.timer:
      self.timer.cancel()
      self.timer = None

  def _send_arp (self, dstip, src_eth, src_ip, send_function):
    r = pkt.arp()
    r.opcode = r.REQUEST
    r.hwdst = pkt.ETHERNET.ETHER_BROADCAST
    r.protodst = dstip
    r.hwsrc = src_eth
    r.protosrc = src_ip
    e = pkt.ethernet(type=pkt.ethernet.ARP_TYPE,
                     src=r.hwsrc,
                     dst=r.hwdst)
    e.payload = r
    log.debug("Sending ARP for %s", dstip)
    send_function(e.pack())

  def rx_arp_reply (self, arp):
    assert arp.opcode == arp.REPLY
    self.rx_arp(arp)

  def rx_arp (self, arp):
    if arp.protosrc not in self.by_ip:
      self.add_entry(mac=arp.hwsrc, ip=arp.protosrc)
    else:
      self.by_ip[arp.protosrc].confirm(arp.hwsrc)

    # Send any pending packets
    for index,(ip,args) in reversed(list(enumerate(self.pending))):
      if ip == arp.protosrc:
        del self.pending[index]
        log.debug("ARP reply allows sending pending packet")
        self.send(*args)
<commit_msg>arp_table: Fill in dst eth for all multicasts<commit_after># Copyright 2017 James McCauley
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Reusable ARP table
"""

import time
from pox.lib.recoco import Timer
from pox.core import core
import pox.lib.packet as pkt

log = core.getLogger()

class ARPEntry (object):
  REFRESH_TIME = 60 # Refresh this often
  MIN_BACKOFF = 1
  MAX_BACKOFF = 32

  def __init__ (self, ip, mac):
    self.ts = time.time()
    self.ip = ip
    self.mac = mac
    if mac is not None:
      self.confirm()
    else:
      self.next_refresh = time.time()
      self.refresh_backoff = self.MIN_BACKOFF

  @property
  def age (self):
    return time.time() - self.ts

  def confirm (self, mac=None):
    self.next_refresh = time.time() + self.REFRESH_TIME
    self.refresh_backoff = self.MIN_BACKOFF
    if mac is not None:
      self.mac = mac

  def maybe_refresh (self):
    """
    if returns True, you should send a refresh
    """
    if self.next_refresh > time.time():
      return False
    self.next_refresh = time.time() + self.refresh_backoff
    self.refresh_backoff *= 2
    if self.refresh_backoff > self.MAX_BACKOFF:
      self.refresh_backoff = self.MAX_BACKOFF
    return True

class ARPTable (object):
  MAX_ENTRIES = 1024
  MAX_PENDING = 5

  MAX_PENDING_TIME = 2

  def __init__ (self):
    self.by_ip = {} # ip -> entry
    self.pending = [] # Packets waiting to be sent (ip,(send_args))
    self.timer = Timer(self.MAX_PENDING_TIME, self._timer_proc)

  def add_entry (self, ip, mac=None):
    """
    Add an entry

    The entry can't already exist.
    It will definitely exist after returning.
    """
    assert ip not in self.by_ip
    if len(self.by_ip) >= self.MAX_ENTRIES:
      # Sloppy, but simple.
      # Get ones with lowest age
      entries = sorted(self.by_ip.values(), key = lambda entry: entry.age)
      del entries[self.MAX_ENTRIES:]
      self.by_ip = {e.mac:e for e in entries}
    new_entry = ARPEntry(ip=ip, mac=mac)
    self.by_ip[ip] = new_entry
    return new_entry

  def send (self, eth_packet, router_ip=None, src_eth=None, src_ip=None,
            send_function=None):
    """
    Try to send a packet

    eth_packet is an ethernet object.
    src_eth is the source for any ARPs sent.
    src_ip is the source for any ARPs sent.
    If the above two are not specified, they are taken from eth_packet.
    send_function is a function which takes raw bytes to send.
    If send_function is unset, it is taken from a send_function attribute.
    """
    if send_function is None: send_function = self.send_function
    ipp = eth_packet.find("ipv4")

    if not ipp and eth_packet.type == eth_packet.IP_TYPE:
      if isinstance(eth_packet.payload, bytes):
        # Hm!  Try harder...
        ipp = pkt.ipv4(raw=eth_packet.payload)

    if not ipp or eth_packet.dst.is_multicast:
      send_function(eth_packet.pack())
      return
    if ipp.dstip == pkt.IPV4.IP_BROADCAST:
      #ipp.dstip = router_ip # Not sure what this was about
      eth_packet.dst = pkt.ETHERNET.ETHER_BROADCAST
      send_function(eth_packet.pack())
      return
    if ipp.dstip.is_multicast:
      eth_packet.dst = ipp.dstip.multicast_ethernet_address
      send_function(eth_packet.pack())
      return

    if src_ip is None: src_ip = ipp.srcip
    if src_eth is None: src_eth = eth_packet.src

    if router_ip is not None: dstip = router_ip
    else:                     dstip = ipp.dstip

    if dstip not in self.by_ip: self.add_entry(dstip)

    e = self.by_ip[dstip]
    if e.maybe_refresh():
      # Send ARP
      self._send_arp(dstip, src_eth, src_ip,
                     send_function)

    if e.mac is not None:
      eth_packet.dst = e.mac
      send_function(eth_packet.pack())
    else:
      if len(self.pending) < self.MAX_PENDING:
        self.pending.append((dstip, (eth_packet, router_ip, src_eth, src_ip,
                                     send_function)))

  def _timer_proc (self):
    # We just blow away all the entries every interval, so on average, they
    # live for half the interval.
    del self.pending[:]

  def __del__ (self):
    if self.timer:
      self.timer.cancel()
      self.timer = None

  def _send_arp (self, dstip, src_eth, src_ip, send_function):
    r = pkt.arp()
    r.opcode = r.REQUEST
    r.hwdst = pkt.ETHERNET.ETHER_BROADCAST
    r.protodst = dstip
    r.hwsrc = src_eth
    r.protosrc = src_ip
    e = pkt.ethernet(type=pkt.ethernet.ARP_TYPE,
                     src=r.hwsrc,
                     dst=r.hwdst)
    e.payload = r
    log.debug("Sending ARP for %s", dstip)
    send_function(e.pack())

  def rx_arp_reply (self, arp):
    assert arp.opcode == arp.REPLY
    self.rx_arp(arp)

  def rx_arp (self, arp):
    if arp.protosrc not in self.by_ip:
      self.add_entry(mac=arp.hwsrc, ip=arp.protosrc)
    else:
      self.by_ip[arp.protosrc].confirm(arp.hwsrc)

    # Send any pending packets
    for index,(ip,args) in reversed(list(enumerate(self.pending))):
      if ip == arp.protosrc:
        del self.pending[index]
        log.debug("ARP reply allows sending pending packet")
        self.send(*args)
<|endoftext|>