<commit_before>"""
Parser for NWS Climate Report text format
"""
import re
import datetime

from pyiem.nws.product import TextProduct

HEADLINE_RE = re.compile(r"\.\.\.THE ([A-Z_\.\-\(\)\/\,\s]+) CLIMATE SUMMARY FOR\s+([A-Z]+\s[0-9]+\s+[0-9]{4})( CORRECTION)?\.\.\.")

class CLIException(Exception):
    """ Exception """
    pass

def trace(val):
    """ This value could be T or M, account for it! """
    if val == 'M' or val == 'MM':
        return None
    if val == 'T':
        return 0.0001
    return float(val)

def trace_r(val):
    """ Convert our value back into meaningful string """
    if val == 0.0001:
        return 'Trace'
    return val

def parse_snowfall(lines, data):
    """ Parse the snowfall data
WEATHER ITEM   OBSERVED TIME   RECORD YEAR NORMAL DEPARTURE LAST
                VALUE   (LST)  VALUE       VALUE  FROM      YEAR
                                                  NORMAL
SNOWFALL (IN)
  YESTERDAY        0.0          MM      MM   0.0    0.0      0.0
  MONTH TO DATE    0.0                       0.0    0.0      0.0
  SINCE JUN 1      0.0                       0.0    0.0      0.0
  SINCE JUL 1      0.0                       0.0    0.0      0.0
  SNOW DEPTH       0
    """
    for linenum, line in enumerate(lines):
        line = (line+" ").replace(" T ", "0.0001")
        tokens = line.split()
        numbers = re.findall("[-+]?\d*\.\d+|\d+", line)
        if len(numbers) == 0:
            continue
        if line.startswith("YESTERDAY") or line.startswith("TODAY"):
            data['snow_today'] = trace(numbers[0])
            if len(tokens) == 7 and tokens[2] != 'MM' and tokens[3] != 'MM':
                data['snow_today_record'] = float(numbers[1])
                data['snow_today_record_years'] = [int(numbers[2]),]
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['snow_today_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        elif line.startswith("MONTH TO DATE"):
            data['snow_month'] = trace(tokens[3])
        elif line.startswith("SINCE JUN 1"):
            data['snow_jun1'] = trace(tokens[3])
        elif line.startswith("SINCE JUL 1"):
            data['snow_jul1'] = trace(tokens[3])
        elif line.startswith("SINCE DEC 1"):
            data['snow_dec1'] = trace(tokens[3])

def parse_precipitation(lines, data):
    """ Parse the precipitation data """
    for linenum, line in enumerate(lines):
        # careful here as if T is only value, the trailing space is stripped
        line = (line+" ").replace(" T ", "0.0001")
        numbers = re.findall("[-+]?\d*\.\d+|\d+", line)
        print numbers
        if line.startswith("YESTERDAY") or line.startswith("TODAY"):
            if len(numbers) == 0:
                continue
            data['precip_today'] = float(numbers[0])
            if len(numbers) == 6:
                data['precip_today_normal'] = float(numbers[3])
                data['precip_today_record'] = float(numbers[1])
                data['precip_today_record_years'] = [int(numbers[2]),]
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['precip_today_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        elif line.startswith("MONTH TO DATE"):
            data['precip_month'] = float(numbers[0])
            if len(numbers) == 4:
                data['precip_month_normal'] = float(numbers[1])
        elif line.startswith("SINCE JAN 1"):
            data['precip_jan1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jan1_normal'] = float(numbers[2])
        elif line.startswith("SINCE JUL 1"):
            data['precip_jul1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jul1_normal'] = float(numbers[2])
        elif line.startswith("SINCE JUN 1"):
            data['precip_jun1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jun1_normal'] = float(numbers[2])
        elif line.startswith("SINCE DEC 1"):
            data['precip_dec1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_dec1_normal'] = float(numbers[2])

def no99(val):
    """ Giveme int val of null! """
    if val == '-99':
        return None
    return int(val)

def parse_temperature(lines, data):
    """ Here we parse a temperature section
WEATHER ITEM   OBSERVED TIME   RECORD YEAR NORMAL DEPARTURE LAST
                VALUE   (LST)  VALUE       VALUE  FROM      YEAR
                                                  NORMAL
..................................................................
TEMPERATURE (F)
 YESTERDAY
  MAXIMUM         89    309 PM 101    1987  85      4       99
  MINIMUM         63    545 AM  51    1898  67     -4       69
  AVERAGE         76                        76      0       84
    """
    for linenum, line in enumerate(lines):
        numbers = re.findall("\-?\d+", line.replace(" MM ", " -99 "))
        if line.startswith("MAXIMUM"):
            data['temperature_maximum'] = no99(numbers[0])
            tokens = re.findall("([0-9]{3,4} [AP]M)", line)
            if len(tokens) == 1:
                data['temperature_maximum_time'] = tokens[0]
            if len(numbers) == 7: # we know this
                data['temperature_maximum_record'] = no99(numbers[2])
                if int(numbers[3]) > 0:
                    data['temperature_maximum_record_years'] = [int(numbers[3]),]
                data['temperature_maximum_normal'] = no99(numbers[4])
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['temperature_maximum_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        if line.startswith("MINIMUM"):
            data['temperature_minimum'] = no99(numbers[0])
            tokens = re.findall("([0-9]{3,4} [AP]M)", line)
            if len(tokens) == 1:
                data['temperature_minimum_time'] = tokens[0]
            if len(numbers) == 7: # we know this
                data['temperature_minimum_record'] = no99(numbers[2])
                if int(numbers[3]) > 0:
                    data['temperature_minimum_record_years'] = [int(numbers[3]),]
                data['temperature_minimum_normal'] = no99(numbers[4])
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['temperature_minimum_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1

class CLIProduct( TextProduct ):
    """
    Represents a Storm Prediction Center Mesoscale Convective Discussion
    """

    def __init__(self, text):
        """ constructor """
        TextProduct.__init__(self, text)
        self.data = None
        self.cli_valid = None
        self.cli_station = None
        if self.wmo[:2] != 'CD':
            print 'Product %s skipped due to wrong header' % (
                                                    self.get_product_id(),)
            return
        self.parse_cli_headline()
        # If we failed above
        if self.cli_valid is not None:
            self.data = self.parse_data()

    def get_jabbers(self, uri, _=None):
        """ Override the jabber message formatter """
        url = "%s?pid=%s" % (uri, self.get_product_id())

        mess = "%s %s Climate Report: High: %s Low: %s Precip: %s Snow: %s %s" % (
                    self.cli_station, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M')), url
                    )
        htmlmess = ("%s %s Climate Report: High: %s "
                    +"Low: %s Precip: %s Snow: %s") % (
                    self.cli_station, url, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M'))
                    )
        tweet = "%s %s Climate: Hi: %s Lo: %s Precip: %s Snow: %s %s" % (
                    self.cli_station, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M')), url
                    )
        res = []
        res.append([mess.replace("0.0001", "Trace"),
                    htmlmess.replace("0.0001", "Trace"), {
                        'channels': self.afos,
                        'twitter': tweet
                }])
        return res

    def parse_data(self):
        """ Actually do the parsing of this silly format """
        data = {}
        pos = self.unixtext.find("TEMPERATURE")
        if pos == -1:
            raise CLIException('Failed to find TEMPERATURE, aborting')

        # Strip extraneous spaces
        meat = "\n".join([l.strip() for l in self.unixtext[pos:].split("\n")])
        sections = meat.split("\n\n")
        for section in sections:
            lines = section.split("\n")
            if lines[0] in ["TEMPERATURE (F)", 'TEMPERATURE']:
                parse_temperature(lines, data)
            elif lines[0] in ['PRECIPITATION (IN)', 'PRECIPITATION']:
                parse_precipitation(lines, data)
            elif lines[0] in ['SNOWFALL (IN)', 'SNOWFALL']:
                parse_snowfall(lines, data)

        return data

    def parse_cli_headline(self):
        """ Figure out when this product is valid for """
        tokens = HEADLINE_RE.findall( self.unixtext.replace("\n", " ") )
        if len(tokens) == 1:
            if len(tokens[0][1].split()[0]) == 3:
                myfmt = '%b %d %Y'
            else:
                myfmt = '%B %d %Y'
            self.cli_valid = datetime.datetime.strptime(tokens[0][1], myfmt)
            self.cli_station = (tokens[0][0]).strip()
        elif len(tokens) > 1:
            raise CLIException("Found two headers in product, unsupported!")
        else:
            # Known sources of bad data...
            if self.source in ['PKMR', 'NSTU', 'PTTP', 'PTKK', 'PTKR']:
                return None
            raise CLIException('Could not find date valid in %s' % (
                                                self.get_product_id(),))

def parser(text, utcnow=None, ugc_provider=None, nwsli_provider=None):
    """ Provide back CLI objects based on the parsing of this text """
    return CLIProduct( text )
<commit_msg>another attempt at getting snowfall right<commit_after>"""
Parser for NWS Climate Report text format
"""
import re
import datetime

from pyiem.nws.product import TextProduct

HEADLINE_RE = re.compile(r"\.\.\.THE ([A-Z_\.\-\(\)\/\,\s]+) CLIMATE SUMMARY FOR\s+([A-Z]+\s[0-9]+\s+[0-9]{4})( CORRECTION)?\.\.\.")

class CLIException(Exception):
    """ Exception """
    pass

def trace(val):
    """ This value could be T or M, account for it! """
    if val == 'M' or val == 'MM':
        return None
    if val == 'T':
        return 0.0001
    return float(val)

def trace_r(val):
    """ Convert our value back into meaningful string """
    if val == 0.0001:
        return 'Trace'
    return val

def parse_snowfall(lines, data):
    """ Parse the snowfall data
WEATHER ITEM   OBSERVED TIME   RECORD YEAR NORMAL DEPARTURE LAST
                VALUE   (LST)  VALUE       VALUE  FROM      YEAR
                                                  NORMAL
SNOWFALL (IN)
  YESTERDAY        0.0          MM      MM   0.0    0.0      0.0
  MONTH TO DATE    0.0                       0.0    0.0      0.0
  SINCE JUN 1      0.0                       0.0    0.0      0.0
  SINCE JUL 1      0.0                       0.0    0.0      0.0
  SNOW DEPTH       0
    """
    for linenum, line in enumerate(lines):
        line = (line+" ").replace(" T ", "0.0001").replace("R ", "  ")
        tokens = line.split()
        numbers = re.findall("[-+]?\d*\.\d+|\d+", line)
        if len(numbers) == 0:
            continue
        if line.startswith("YESTERDAY") or line.startswith("TODAY"):
            data['snow_today'] = trace(tokens[1])
            if len(tokens) == 7 and tokens[2] != 'MM' and tokens[3] != 'MM':
                data['snow_today_record'] = trace(tokens[2])
                data['snow_today_record_years'] = [int(tokens[3]),]
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['snow_today_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        elif line.startswith("MONTH TO DATE"):
            data['snow_month'] = trace(tokens[3])
        elif line.startswith("SINCE JUN 1"):
            data['snow_jun1'] = trace(tokens[3])
        elif line.startswith("SINCE JUL 1"):
            data['snow_jul1'] = trace(tokens[3])
        elif line.startswith("SINCE DEC 1"):
            data['snow_dec1'] = trace(tokens[3])

def parse_precipitation(lines, data):
    """ Parse the precipitation data """
    for linenum, line in enumerate(lines):
        # careful here as if T is only value, the trailing space is stripped
        line = (line+" ").replace(" T ", "0.0001")
        numbers = re.findall("[-+]?\d*\.\d+|\d+", line)
        if line.startswith("YESTERDAY") or line.startswith("TODAY"):
            if len(numbers) == 0:
                continue
            data['precip_today'] = float(numbers[0])
            if len(numbers) == 6:
                data['precip_today_normal'] = float(numbers[3])
                data['precip_today_record'] = float(numbers[1])
                data['precip_today_record_years'] = [int(numbers[2]),]
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['precip_today_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        elif line.startswith("MONTH TO DATE"):
            data['precip_month'] = float(numbers[0])
            if len(numbers) == 4:
                data['precip_month_normal'] = float(numbers[1])
        elif line.startswith("SINCE JAN 1"):
            data['precip_jan1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jan1_normal'] = float(numbers[2])
        elif line.startswith("SINCE JUL 1"):
            data['precip_jul1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jul1_normal'] = float(numbers[2])
        elif line.startswith("SINCE JUN 1"):
            data['precip_jun1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_jun1_normal'] = float(numbers[2])
        elif line.startswith("SINCE DEC 1"):
            data['precip_dec1'] = float(numbers[1])
            if len(numbers) == 4:
                data['precip_dec1_normal'] = float(numbers[2])

def no99(val):
    """ Giveme int val of null! """
    if val == '-99':
        return None
    return int(val)

def parse_temperature(lines, data):
    """ Here we parse a temperature section
WEATHER ITEM   OBSERVED TIME   RECORD YEAR NORMAL DEPARTURE LAST
                VALUE   (LST)  VALUE       VALUE  FROM      YEAR
                                                  NORMAL
..................................................................
TEMPERATURE (F)
 YESTERDAY
  MAXIMUM         89    309 PM 101    1987  85      4       99
  MINIMUM         63    545 AM  51    1898  67     -4       69
  AVERAGE         76                        76      0       84
    """
    for linenum, line in enumerate(lines):
        numbers = re.findall("\-?\d+", line.replace(" MM ", " -99 "))
        if line.startswith("MAXIMUM"):
            data['temperature_maximum'] = no99(numbers[0])
            tokens = re.findall("([0-9]{3,4} [AP]M)", line)
            if len(tokens) == 1:
                data['temperature_maximum_time'] = tokens[0]
            if len(numbers) == 7: # we know this
                data['temperature_maximum_record'] = no99(numbers[2])
                if int(numbers[3]) > 0:
                    data['temperature_maximum_record_years'] = [int(numbers[3]),]
                data['temperature_maximum_normal'] = no99(numbers[4])
                # Check next line(s) for more years
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['temperature_maximum_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1
        if line.startswith("MINIMUM"):
            data['temperature_minimum'] = no99(numbers[0])
            tokens = re.findall("([0-9]{3,4} [AP]M)", line)
            if len(tokens) == 1:
                data['temperature_minimum_time'] = tokens[0]
            if len(numbers) == 7: # we know this
                data['temperature_minimum_record'] = no99(numbers[2])
                if int(numbers[3]) > 0:
                    data['temperature_minimum_record_years'] = [int(numbers[3]),]
                data['temperature_minimum_normal'] = no99(numbers[4])
                while ((linenum+1)<len(lines) and
                       len(lines[linenum+1].strip()) == 4):
                    data['temperature_minimum_record_years'].append(
                                                    int(lines[linenum+1]))
                    linenum += 1

class CLIProduct( TextProduct ):
    """
    Represents a Storm Prediction Center Mesoscale Convective Discussion
    """

    def __init__(self, text):
        """ constructor """
        TextProduct.__init__(self, text)
        self.data = None
        self.cli_valid = None
        self.cli_station = None
        if self.wmo[:2] != 'CD':
            print 'Product %s skipped due to wrong header' % (
                                                    self.get_product_id(),)
            return
        self.parse_cli_headline()
        # If we failed above
        if self.cli_valid is not None:
            self.data = self.parse_data()

    def get_jabbers(self, uri, _=None):
        """ Override the jabber message formatter """
        url = "%s?pid=%s" % (uri, self.get_product_id())

        mess = "%s %s Climate Report: High: %s Low: %s Precip: %s Snow: %s %s" % (
                    self.cli_station, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M')), url
                    )
        htmlmess = ("%s %s Climate Report: High: %s "
                    +"Low: %s Precip: %s Snow: %s") % (
                    self.cli_station, url, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M'))
                    )
        tweet = "%s %s Climate: Hi: %s Lo: %s Precip: %s Snow: %s %s" % (
                    self.cli_station, self.cli_valid.strftime("%b %-d"),
                    self.data.get('temperature_maximum', 'M'),
                    self.data.get('temperature_minimum', 'M'),
                    trace_r(self.data.get('precip_today', 'M')),
                    trace_r(self.data.get('snow_today', 'M')), url
                    )
        res = []
        res.append([mess.replace("0.0001", "Trace"),
                    htmlmess.replace("0.0001", "Trace"), {
                        'channels': self.afos,
                        'twitter': tweet
                }])
        return res

    def parse_data(self):
        """ Actually do the parsing of this silly format """
        data = {}
        pos = self.unixtext.find("TEMPERATURE")
        if pos == -1:
            raise CLIException('Failed to find TEMPERATURE, aborting')

        # Strip extraneous spaces
        meat = "\n".join([l.strip() for l in self.unixtext[pos:].split("\n")])
        sections = meat.split("\n\n")
        for section in sections:
            lines = section.split("\n")
            if lines[0] in ["TEMPERATURE (F)", 'TEMPERATURE']:
                parse_temperature(lines, data)
            elif lines[0] in ['PRECIPITATION (IN)', 'PRECIPITATION']:
                parse_precipitation(lines, data)
            elif lines[0] in ['SNOWFALL (IN)', 'SNOWFALL']:
                parse_snowfall(lines, data)

        return data

    def parse_cli_headline(self):
        """ Figure out when this product is valid for """
        tokens = HEADLINE_RE.findall( self.unixtext.replace("\n", " ") )
        if len(tokens) == 1:
            if len(tokens[0][1].split()[0]) == 3:
                myfmt = '%b %d %Y'
            else:
                myfmt = '%B %d %Y'
            self.cli_valid = datetime.datetime.strptime(tokens[0][1], myfmt)
            self.cli_station = (tokens[0][0]).strip()
        elif len(tokens) > 1:
            raise CLIException("Found two headers in product, unsupported!")
        else:
            # Known sources of bad data...
            if self.source in ['PKMR', 'NSTU', 'PTTP', 'PTKK', 'PTKR']:
                return None
            raise CLIException('Could not find date valid in %s' % (
                                                self.get_product_id(),))

def parser(text, utcnow=None, ugc_provider=None, nwsli_provider=None):
    """ Provide back CLI objects based on the parsing of this text """
    return CLIProduct( text )
<|endoftext|>


<commit_before>'''
Manage git repositories and configuration.
'''

from __future__ import unicode_literals

import re

from pyinfra.api import operation, OperationError

from . import files, ssh
from .util.files import chown

@operation(pipeline_facts={
    'git_config': 'repo',
})
def config(
    state, host, key, value,
    repo=None,
):
    '''
    Manage git config for a repository or globally.

    + key: the key of the config to ensure
    + value: the value this key should have
    + repo: specify the git repo path to edit local config (defaults to global)

    Example:

    .. code:: python

        git.config(
            name='Ensure user name is set for a repo',
            key='user.name',
            value='Anon E. Mouse',
            repo='/usr/local/src/pyinfra',
        )

    '''

    existing_config = host.fact.git_config(repo)

    if not existing_config or existing_config.get(key) != value:
        if repo is None:
            yield 'git config --global {0} "{1}"'.format(key, value)
        else:
            yield 'cd {0} && git config --local {1} "{2}"'.format(repo, key, value)

@operation(pipeline_facts={
    'git_branch': 'target',
})
def repo(
    state, host, source, target,
    branch='master', pull=True, rebase=False,
    user=None, group=None, ssh_keyscan=False,
    update_submodules=False, recursive_submodules=False,
):
    '''
    Clone/pull git repositories.

    + source: the git source URL
    + target: target directory to clone to
    + branch: branch to pull/checkout
    + pull: pull any changes for the branch
    + rebase: when pulling, use ``--rebase``
    + user: chown files to this user after
    + group: chown files to this group after
    + ssh_keyscan: keyscan the remote host if not in known_hosts before clone/pull
    + update_submodules: update any git submodules
    + recursive_submodules: update git submodules recursively

    Example:

    .. code:: python

        git.repo(
            name='Clone repo',
            src='https://github.com/Fizzadar/pyinfra.git',
            dest='/usr/local/src/pyinfra',
        )
    '''

    # Ensure our target directory exists
    yield files.directory(state, host, target)

    # Do we need to scan for the remote host key?
    if ssh_keyscan:
        # Attempt to parse the domain from the git repository
        domain = re.match(r'^[a-zA-Z0-9]+@([0-9a-zA-Z\.\-]+)', source)

        if domain:
            yield ssh.keyscan(state, host, domain.group(1))
        else:
            raise OperationError(
                'Could not parse domain (to SSH keyscan) from: {0}'.format(source),
            )

    # Store git commands for directory prefix
    git_commands = []
    is_repo = host.fact.directory('/'.join((target, '.git')))

    # Cloning new repo?
    if not is_repo:
        git_commands.append('clone {0} --branch {1} .'.format(source, branch))

    # Ensuring existing repo
    else:
        current_branch = host.fact.git_branch(target)
        if current_branch != branch:
            git_commands.append('fetch')  # fetch to ensure we have the branch locally
            git_commands.append('checkout {0}'.format(branch))

        if pull:
            if rebase:
                git_commands.append('pull --rebase')
            else:
                git_commands.append('pull')

    if update_submodules:
        if recursive_submodules:
            git_commands.append('submodule update --init --recursive')
        else:
            git_commands.append('submodule update --init')

    # Attach prefixes for directory
    command_prefix = 'cd {0} && git'.format(target)
    git_commands = [
        '{0} {1}'.format(command_prefix, command)
        for command in git_commands
    ]

    for cmd in git_commands:
        yield cmd

    # Apply any user or group
    if user or group:
        yield chown(target, user, group, recursive=True)
<commit_msg>Update `git.repo` arguments:<commit_after>'''
Manage git repositories and configuration.
'''

from __future__ import unicode_literals

import re

from pyinfra.api import operation, OperationError

from . import files, ssh
from .util.files import chown

@operation(pipeline_facts={
    'git_config': 'repo',
})
def config(
    state, host, key, value,
    repo=None,
):
    '''
    Manage git config for a repository or globally.

    + key: the key of the config to ensure
    + value: the value this key should have
    + repo: specify the git repo path to edit local config (defaults to global)

    Example:

    .. code:: python

        git.config(
            name='Ensure user name is set for a repo',
            key='user.name',
            value='Anon E. Mouse',
            repo='/usr/local/src/pyinfra',
        )

    '''

    existing_config = host.fact.git_config(repo)

    if not existing_config or existing_config.get(key) != value:
        if repo is None:
            yield 'git config --global {0} "{1}"'.format(key, value)
        else:
            yield 'cd {0} && git config --local {1} "{2}"'.format(repo, key, value)

@operation(pipeline_facts={
    'git_branch': 'target',
})
def repo(
    state, host, src, dest,
    branch='master', pull=True, rebase=False,
    user=None, group=None, ssh_keyscan=False,
    update_submodules=False, recursive_submodules=False,
):
    '''
    Clone/pull git repositories.

    + src: the git source URL
    + dest: directory to clone to
    + branch: branch to pull/checkout
    + pull: pull any changes for the branch
    + rebase: when pulling, use ``--rebase``
    + user: chown files to this user after
    + group: chown files to this group after
    + ssh_keyscan: keyscan the remote host if not in known_hosts before clone/pull
    + update_submodules: update any git submodules
    + recursive_submodules: update git submodules recursively

    Example:

    .. code:: python

        git.repo(
            name='Clone repo',
            src='https://github.com/Fizzadar/pyinfra.git',
            dest='/usr/local/src/pyinfra',
        )
    '''

    # Ensure our target directory exists
    yield files.directory(state, host, dest)

    # Do we need to scan for the remote host key?
    if ssh_keyscan:
        # Attempt to parse the domain from the git repository
        domain = re.match(r'^[a-zA-Z0-9]+@([0-9a-zA-Z\.\-]+)', src)

        if domain:
            yield ssh.keyscan(state, host, domain.group(1))
        else:
            raise OperationError(
                'Could not parse domain (to SSH keyscan) from: {0}'.format(src),
            )

    # Store git commands for directory prefix
    git_commands = []
    is_repo = host.fact.directory('/'.join((dest, '.git')))

    # Cloning new repo?
    if not is_repo:
        git_commands.append('clone {0} --branch {1} .'.format(src, branch))

    # Ensuring existing repo
    else:
        current_branch = host.fact.git_branch(dest)
        if current_branch != branch:
            git_commands.append('fetch')  # fetch to ensure we have the branch locally
            git_commands.append('checkout {0}'.format(branch))

        if pull:
            if rebase:
                git_commands.append('pull --rebase')
            else:
                git_commands.append('pull')

    if update_submodules:
        if recursive_submodules:
            git_commands.append('submodule update --init --recursive')
        else:
            git_commands.append('submodule update --init')

    # Attach prefixes for directory
    command_prefix = 'cd {0} && git'.format(dest)
    git_commands = [
        '{0} {1}'.format(command_prefix, command)
        for command in git_commands
    ]

    for cmd in git_commands:
        yield cmd

    # Apply any user or group
    if user or group:
        yield chown(dest, user, group, recursive=True)
<|endoftext|>


<commit_before>'''
Manage git repositories and configuration.
'''

from __future__ import unicode_literals

import re

from pyinfra import logger
from pyinfra.api import operation, OperationError

from . import files, ssh
from .util.files import chown

@operation(pipeline_facts={
    'git_config': 'repo',
})
def config(
    key, value,
    repo=None,
    state=None, host=None,
):
    '''
    Manage git config for a repository or globally.

    + key: the key of the config to ensure
    + value: the value this key should have
    + repo: specify the git repo path to edit local config (defaults to global)

    Example:

    .. code:: python

        git.config(
            name='Ensure user name is set for a repo',
            key='user.name',
            value='Anon E. Mouse',
            repo='/usr/local/src/pyinfra',
        )

    '''

    existing_config = {}

    if not repo:
        existing_config = host.fact.git_config()

    # Only get the config if the repo exists at this stage
    elif host.fact.directory('/'.join((repo, '.git'))):
        existing_config = host.fact.git_config(repo)

    if existing_config.get(key) != value:
        if repo is None:
            yield 'git config --global {0} "{1}"'.format(key, value)
        else:
            yield 'cd {0} && git config --local {1} "{2}"'.format(repo, key, value)

        existing_config[key] = value

    else:
        host.noop('git config {0} is set to {1}'.format(key, value))

@operation(pipeline_facts={
    'git_branch': 'target',
})
def repo(
    src, dest,
    branch='master', pull=True, rebase=False,
    user=None, group=None, ssh_keyscan=False,
    update_submodules=False, recursive_submodules=False,
    state=None, host=None,
):
    '''
    Clone/pull git repositories.

    + src: the git source URL
    + dest: directory to clone to
    + branch: branch to pull/checkout
    + pull: pull any changes for the branch
    + rebase: when pulling, use ``--rebase``
    + user: chown files to this user after
    + group: chown files to this group after
    + ssh_keyscan: keyscan the remote host if not in known_hosts before clone/pull
    + update_submodules: update any git submodules
    + recursive_submodules: update git submodules recursively

    Example:

    .. code:: python

        git.repo(
            name='Clone repo',
            src='https://github.com/Fizzadar/pyinfra.git',
            dest='/usr/local/src/pyinfra',
        )
    '''

    # Ensure our target directory exists
    yield files.directory(dest, state=state, host=host)

    # Do we need to scan for the remote host key?
    if ssh_keyscan:
        # Attempt to parse the domain from the git repository
        domain = re.match(r'^[a-zA-Z0-9]+@([0-9a-zA-Z\.\-]+)', src)

        if domain:
            yield ssh.keyscan(domain.group(1), state=state, host=host)
        else:
            raise OperationError(
                'Could not parse domain (to SSH keyscan) from: {0}'.format(src),
            )

    # Store git commands for directory prefix
    git_commands = []
    is_repo = host.fact.directory('/'.join((dest, '.git')))

    # Cloning new repo?
    if not is_repo:
        git_commands.append('clone {0} --branch {1} .'.format(src, branch))

    # Ensuring existing repo
    else:
        current_branch = host.fact.git_branch(dest)
        if current_branch != branch:
            git_commands.append('fetch')  # fetch to ensure we have the branch locally
            git_commands.append('checkout {0}'.format(branch))

        if pull:
            if rebase:
                git_commands.append('pull --rebase')
            else:
                git_commands.append('pull')

    if update_submodules:
        if recursive_submodules:
            git_commands.append('submodule update --init --recursive')
        else:
            git_commands.append('submodule update --init')

    # Attach prefixes for directory
    command_prefix = 'cd {0} && git'.format(dest)
    git_commands = [
        '{0} {1}'.format(command_prefix, command)
        for command in git_commands
    ]

    for cmd in git_commands:
        yield cmd

    # Apply any user or group
    if user or group:
        yield chown(dest, user, group, recursive=True)

@operation()
def worktree(
    worktree,
    repo=None, branch=None, create_branch=False, detached=False,
    new_branch=None, commitish=None,
    pull=True, rebase=False, from_remote_branch=None,
    present=True, assume_repo_exists=False, force=False,
    user=None, group=None, state=None, host=None,
):
    '''
    Manage git worktrees.

    + worktree: git working tree directory
    + repo: git main repository directory
    + detached: create a working tree with a detached HEAD
    + branch: (deprecated)
    + create_branch: (deprecated)
    + new_branch: local branch name created at the same time than the worktree
    + commitish: from which git commit, branch, ... the worktree is created
    + pull: pull any changes from a remote branch if set
    + rebase: when pulling, use ``--rebase``
    + from_remote_branch: a 2-tuple ``(remote, branch)`` that identifies a remote branch
    + present: whether the working tree should exist
    + assume_repo_exists: whether to assume the main repo exists
    + force: remove unclean working tree if should not exist
    + user: chown files to this user after
    + group: chown files to this group after

    Example:

    .. code:: python

        git.worktree(
            name='Create a worktree from the current repo `HEAD`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix'
        )

        git.worktree(
            name='Create a worktree from the commit `4e091aa0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='4e091aa0'
        )

        git.worktree(
            name='Create a worktree with a new local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
        )

        git.worktree(
            name='Create a worktree from the commit 4e091aa0 with the new local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
            commitish='4e091aa0'
        )

        git.worktree(
            name='Create a worktree with a detached `HEAD`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            detached=True,
        )

        git.worktree(
            name='Create a worktree with a detached `HEAD` from commit `4e091aa0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='4e091aa0',
            detached=True,
        )

        git.worktree(
            name='Create a worktree from the existing local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='v1.0'
        )

        git.worktree(
            name='Create a worktree with a new branch `v1.0` that tracks `origin/v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
            commitish='v1.0'
        )

        git.worktree(
            name='Pull an existing worktree already linked to a tracking branch',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix'
        )

        git.worktree(
            name='Pull an existing worktree from a specific remote branch',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            from_remote_branch=('origin', 'master')
        )

        git.worktree(
            name='Remove a worktree',
            worktree='/usr/local/src/pyinfra/hotfix',
            present=False,
        )

        git.worktree(
            name='Remove an unclean worktree',
            worktree='/usr/local/src/pyinfra/hotfix',
            present=False,
            force=True,
        )
    '''

    # handle deprecated arguments.
    #  1) old api: branch=xxx, create_branch=True => git worktree add -b xxx <worktree_path>
    #     new api: new_branch=xxx
    #
    #  2) old api: branch=xxx, create_branch=False => git worktree add <worktree_path> xxx
    #     new api: commitish=xxx
    if branch:
        logger.warning(
            'The `branch` and `create_branch` arguments are deprecated. '
            'Please use `branch` and `commitish`.',
        )

        if create_branch:
            if new_branch:
                raise OperationError(
                    'The deprecated arguments `branch` and `create_branch` are not compatible with'
                    ' the new `branch` argument. Please use the new `branch` argument only.',
                )
            else:
                new_branch = branch
        else:
            if commitish:
                raise OperationError(
                    'The deprecated arguments `branch` and `create_branch` are not compatible with'
                    ' the new `commitish` argument. Please use the new `commitish` argument only.',
                )
            else:
                commitish = branch

    # Doesn't exist & we want it
    if not host.fact.directory(worktree) and present:

        # be sure that `repo` is a GIT repository
        if not assume_repo_exists and not host.fact.directory('/'.join((repo, '.git'))):
            raise OperationError(
                'The following folder is not a valid GIT repository : {0}'.format(repo),
            )

        command_parts = ['cd {0} && git worktree add'.format(repo)]

        if new_branch:
            command_parts.append('-b {0}'.format(new_branch))
        elif detached:
            command_parts.append('--detach')

        command_parts.append(worktree)

        if commitish:
            command_parts.append(commitish)

        yield ' '.join(command_parts)

        # Apply any user or group
        if user or group:
            yield chown(worktree, user, group, recursive=True)

    # It exists and we don't want it
    elif host.fact.directory(worktree) and not present:

        command = 'cd {0} && git worktree remove .'.format(worktree)

        if force:
            command += ' --force'

        yield command

    # It exists and we still want it => pull/rebase it
    elif host.fact.directory(worktree) and present:

        # pull the worktree only if it's already linked to a tracking branch or
        # if a remote branch is set
        if host.fact.git_tracking_branch(worktree) or from_remote_branch:
            command = 'cd {0} && git pull'.format(worktree)

            if rebase:
                command += ' --rebase'

            if from_remote_branch:
                if len(from_remote_branch) != 2 or type(from_remote_branch) not in (tuple, list):
                    raise OperationError(
                        'The remote branch must be a 2-tuple (remote, branch) such as '
                        '("origin", "master")',
                    )
                command += ' {0} {1}'.format(*from_remote_branch)

            yield command

@operation
def bare_repo(
    path,
    user=None,
    group=None,
    present=True,
    state=None,
    host=None,
):
    '''
    Create bare git repositories.

    + path: path to the folder
    + present: whether the bare repository should exist
    + user: chown files to this user after
    + group: chown files to this group after

    Example:

    .. code:: python

        git.bare_repo(
            name='Create bare repo',
            path='/home/git/test.git',
        )
    '''

    yield files.directory(path, state=state, host=host, present=present)

    # Ensure our target directory exists
    if present:
        yield 'git init --bare {0}'.format(path)

    # Apply any user or group
    if user or group:
        yield chown(path, user, group, recursive=True)
<commit_msg>Support `branch=None` in `git.repo` operation.<commit_after>'''
Manage git repositories and configuration.
'''

from __future__ import unicode_literals

import re

from pyinfra import logger
from pyinfra.api import operation, OperationError

from . import files, ssh
from .util.files import chown

@operation(pipeline_facts={
    'git_config': 'repo',
})
def config(
    key, value,
    repo=None,
    state=None, host=None,
):
    '''
    Manage git config for a repository or globally.

    + key: the key of the config to ensure
    + value: the value this key should have
    + repo: specify the git repo path to edit local config (defaults to global)

    Example:

    .. code:: python

        git.config(
            name='Ensure user name is set for a repo',
            key='user.name',
            value='Anon E. Mouse',
            repo='/usr/local/src/pyinfra',
        )

    '''

    existing_config = {}

    if not repo:
        existing_config = host.fact.git_config()

    # Only get the config if the repo exists at this stage
    elif host.fact.directory('/'.join((repo, '.git'))):
        existing_config = host.fact.git_config(repo)

    if existing_config.get(key) != value:
        if repo is None:
            yield 'git config --global {0} "{1}"'.format(key, value)
        else:
            yield 'cd {0} && git config --local {1} "{2}"'.format(repo, key, value)

        existing_config[key] = value

    else:
        host.noop('git config {0} is set to {1}'.format(key, value))

@operation(pipeline_facts={
    'git_branch': 'target',
})
def repo(
    src, dest,
    branch='master',  # TODO: change the default to None
    pull=True, rebase=False,
    user=None, group=None, ssh_keyscan=False,
    update_submodules=False, recursive_submodules=False,
    state=None, host=None,
):
    '''
    Clone/pull git repositories.

    + src: the git source URL
    + dest: directory to clone to
    + branch: branch to pull/checkout
    + pull: pull any changes for the branch
    + rebase: when pulling, use ``--rebase``
    + user: chown files to this user after
    + group: chown files to this group after
    + ssh_keyscan: keyscan the remote host if not in known_hosts before clone/pull
    + update_submodules: update any git submodules
    + recursive_submodules: update git submodules recursively

    Example:

    .. code:: python

        git.repo(
            name='Clone repo',
            src='https://github.com/Fizzadar/pyinfra.git',
            dest='/usr/local/src/pyinfra',
        )
    '''

    # Ensure our target directory exists
    yield files.directory(dest, state=state, host=host)

    # Do we need to scan for the remote host key?
    if ssh_keyscan:
        # Attempt to parse the domain from the git repository
        domain = re.match(r'^[a-zA-Z0-9]+@([0-9a-zA-Z\.\-]+)', src)

        if domain:
            yield ssh.keyscan(domain.group(1), state=state, host=host)
        else:
            raise OperationError(
                'Could not parse domain (to SSH keyscan) from: {0}'.format(src),
            )

    # Store git commands for directory prefix
    git_commands = []
    is_repo = host.fact.directory('/'.join((dest, '.git')))

    # Cloning new repo?
    if not is_repo:
        if branch:
            git_commands.append('clone {0} --branch {1} .'.format(src, branch))
        else:
            git_commands.append('clone {0} .'.format(src))

    # Ensuring existing repo
    else:
        current_branch = host.fact.git_branch(dest)
        if current_branch != branch:
            git_commands.append('fetch')  # fetch to ensure we have the branch locally
            git_commands.append('checkout {0}'.format(branch))

        if pull:
            if rebase:
                git_commands.append('pull --rebase')
            else:
                git_commands.append('pull')

    if update_submodules:
        if recursive_submodules:
            git_commands.append('submodule update --init --recursive')
        else:
            git_commands.append('submodule update --init')

    # Attach prefixes for directory
    command_prefix = 'cd {0} && git'.format(dest)
    git_commands = [
        '{0} {1}'.format(command_prefix, command)
        for command in git_commands
    ]

    for cmd in git_commands:
        yield cmd

    # Apply any user or group
    if user or group:
        yield chown(dest, user, group, recursive=True)

@operation()
def worktree(
    worktree,
    repo=None, branch=None, create_branch=False, detached=False,
    new_branch=None, commitish=None,
    pull=True, rebase=False, from_remote_branch=None,
    present=True, assume_repo_exists=False, force=False,
    user=None, group=None, state=None, host=None,
):
    '''
    Manage git worktrees.

    + worktree: git working tree directory
    + repo: git main repository directory
    + detached: create a working tree with a detached HEAD
    + branch: (deprecated)
    + create_branch: (deprecated)
    + new_branch: local branch name created at the same time than the worktree
    + commitish: from which git commit, branch, ... the worktree is created
    + pull: pull any changes from a remote branch if set
    + rebase: when pulling, use ``--rebase``
    + from_remote_branch: a 2-tuple ``(remote, branch)`` that identifies a remote branch
    + present: whether the working tree should exist
    + assume_repo_exists: whether to assume the main repo exists
    + force: remove unclean working tree if should not exist
    + user: chown files to this user after
    + group: chown files to this group after

    Example:

    .. code:: python

        git.worktree(
            name='Create a worktree from the current repo `HEAD`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix'
        )

        git.worktree(
            name='Create a worktree from the commit `4e091aa0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='4e091aa0'
        )

        git.worktree(
            name='Create a worktree with a new local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
        )

        git.worktree(
            name='Create a worktree from the commit 4e091aa0 with the new local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
            commitish='4e091aa0'
        )

        git.worktree(
            name='Create a worktree with a detached `HEAD`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            detached=True,
        )

        git.worktree(
            name='Create a worktree with a detached `HEAD` from commit `4e091aa0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='4e091aa0',
            detached=True,
        )

        git.worktree(
            name='Create a worktree from the existing local branch `v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            commitish='v1.0'
        )

        git.worktree(
            name='Create a worktree with a new branch `v1.0` that tracks `origin/v1.0`',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            new_branch='v1.0',
            commitish='v1.0'
        )

        git.worktree(
            name='Pull an existing worktree already linked to a tracking branch',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix'
        )

        git.worktree(
            name='Pull an existing worktree from a specific remote branch',
            repo='/usr/local/src/pyinfra/master',
            worktree='/usr/local/src/pyinfra/hotfix',
            from_remote_branch=('origin', 'master')
        )

        git.worktree(
            name='Remove a worktree',
            worktree='/usr/local/src/pyinfra/hotfix',
            present=False,
        )

        git.worktree(
            name='Remove an unclean worktree',
            worktree='/usr/local/src/pyinfra/hotfix',
            present=False,
            force=True,
        )
    '''

    # handle deprecated arguments.
    #  1) old api: branch=xxx, create_branch=True => git worktree add -b xxx <worktree_path>
    #     new api: new_branch=xxx
    #
    #  2) old api: branch=xxx, create_branch=False => git worktree add <worktree_path> xxx
    #     new api: commitish=xxx
    if branch:
        logger.warning(
            'The `branch` and `create_branch` arguments are deprecated. '
            'Please use `branch` and `commitish`.',
        )

        if create_branch:
            if new_branch:
                raise OperationError(
                    'The deprecated arguments `branch` and `create_branch` are not compatible with'
                    ' the new `branch` argument. Please use the new `branch` argument only.',
                )
            else:
                new_branch = branch
        else:
            if commitish:
                raise OperationError(
                    'The deprecated arguments `branch` and `create_branch` are not compatible with'
                    ' the new `commitish` argument. Please use the new `commitish` argument only.',
                )
            else:
                commitish = branch

    # Doesn't exist & we want it
    if not host.fact.directory(worktree) and present:

        # be sure that `repo` is a GIT repository
        if not assume_repo_exists and not host.fact.directory('/'.join((repo, '.git'))):
            raise OperationError(
                'The following folder is not a valid GIT repository : {0}'.format(repo),
            )

        command_parts = ['cd {0} && git worktree add'.format(repo)]

        if new_branch:
            command_parts.append('-b {0}'.format(new_branch))
        elif detached:
            command_parts.append('--detach')

        command_parts.append(worktree)

        if commitish:
            command_parts.append(commitish)

        yield ' '.join(command_parts)

        # Apply any user or group
        if user or group:
            yield chown(worktree, user, group, recursive=True)

    # It exists and we don't want it
    elif host.fact.directory(worktree) and not present:

        command = 'cd {0} && git worktree remove .'.format(worktree)

        if force:
            command += ' --force'

        yield command

    # It exists and we still want it => pull/rebase it
    elif host.fact.directory(worktree) and present:

        # pull the worktree only if it's already linked to a tracking branch or
        # if a remote branch is set
        if host.fact.git_tracking_branch(worktree) or from_remote_branch:
            command = 'cd {0} && git pull'.format(worktree)

            if rebase:
                command += ' --rebase'

            if from_remote_branch:
                if len(from_remote_branch) != 2 or type(from_remote_branch) not in (tuple, list):
                    raise OperationError(
                        'The remote branch must be a 2-tuple (remote, branch) such as '
                        '("origin", "master")',
                    )
                command += ' {0} {1}'.format(*from_remote_branch)

            yield command

@operation
def bare_repo(
    path,
    user=None,
    group=None,
    present=True,
    state=None,
    host=None,
):
    '''
    Create bare git repositories.

    + path: path to the folder
    + present: whether the bare repository should exist
    + user: chown files to this user after
    + group: chown files to this group after

    Example:

    .. code:: python

        git.bare_repo(
            name='Create bare repo',
            path='/home/git/test.git',
        )
    '''

    yield files.directory(path, state=state, host=host, present=present)

    # Ensure our target directory exists
    if present:
        yield 'git init --bare {0}'.format(path)

    # Apply any user or group
    if user or group:
        yield chown(path, user, group, recursive=True)
<|endoftext|>


<commit_before>"""
An interface to the small NORB dataset. Unlike ./norb_small.py, this reads the
original NORB file format, not the LISA lab's .npy version.

Currently only supports the Small NORB Dataset.

Download the dataset from:
http://www.cs.nyu.edu/~ylclab/data/norb-v1.0-small/

NORB dataset(s) by Fu Jie Huang and Yann LeCun.
"""

__authors__ = "Guillaume Desjardins and Matthew Koichi Grimes"
__copyright__ = "Copyright 2010-2014, Universite de Montreal"
__credits__ = __authors__.split(" and ")
__license__ = "3-clause BSD"
__maintainer__ = "Matthew Koichi Grimes"
__email__ = "mkg alum mit edu (@..)"

import os, gzip, bz2, warnings
import numpy, theano
from pylearn2.datasets import dense_design_matrix
from pylearn2.space import VectorSpace, Conv2DSpace, CompositeSpace

class SmallNORB(dense_design_matrix.DenseDesignMatrix):
    """
    An interface to the small NORB dataset.

    If instantiated with default arguments, target labels are integers
    representing categories, which can be looked up using

      category_name = SmallNORB.get_category(label).

    If instantiated with multi_target=True, labels are vectors of indices
    representing:

      [ category, instance, elevation, azimuth, lighting ]

    Like with category, there are class methods that map these ints to their
    actual values, e.g:

      category = SmallNORB.get_category(label[0])
      elevation = SmallNORB.get_elevation_degrees(label[2])
    """

    # Actual image shape may change, e.g. after being preprocessed by
    # datasets.preprocessing.Downsample
    original_image_shape = (96, 96)

    _categories = ['animal',  # four-legged animal
                   'human',  # human figure
                   'airplane',
                   'truck',
                   'car']

    @classmethod
    def get_category(cls, scalar_label):
        return cls._categories[int(scalar_label)]

    @classmethod
    def get_elevation_degrees(cls, scalar_label):
        scalar_label = int(scalar_label)
        assert scalar_label >= 0
        assert scalar_label < 9
        return 30 + 5 * scalar_label

    @classmethod
    def get_azimuth_degrees(cls, scalar_label):
        scalar_label = int(scalar_label)
        assert scalar_label >= 0
        assert scalar_label <= 34
        assert (scalar_label % 2) == 0
        return scalar_label * 10

    # Maps azimuth labels (ints) to their actual values, in degrees.
    azimuth_degrees = numpy.arange(0, 341, 20)

    # Maps a label type to its index within a label vector.
    label_type_to_index = {'category': 0,
                           'instance': 1,
                           'elevation': 2,
                           'azimuth': 3,
                           'lighting': 4}

    # Number of labels, for each label type.
    num_labels_by_type = (len(_categories),
                          10,  # instances
                          9,   # elevations
                          18,  # azimuths
                          6)   # lighting

    # [mkg] Dropped support for the 'center' argument for now. In Pylearn 1, it
    # shifted the pixel values from [0:255] by subtracting 127.5. Seems like a
    # form of preprocessing, which might be better implemented separately using
    # the Preprocess class.
    def __init__(self, which_set, multi_target=False):
        """
        :param which_set: one of ['train', 'test'] :param multi_target: If
        True, each label is an integer labeling the image catergory. If False,
        each label is a vector: [category, instance, lighting, elevation,
        azimuth]. All labels are given as integers. Use the categories,
        elevation_degrees, and azimuth_degrees arrays to map from these
        integers to actual values.

        :param multi_target: If False, labels will be integers indicating
        object category. If True, labels will be vectors of integers,
        indicating [ category, instance, elevation, azimuth, lighting ].
        """

        assert which_set in ['train', 'test']

        self.which_set = which_set

        X = SmallNORB.load(which_set, 'dat')

        # Casts to the GPU-supported float type, using theano._asarray(), a
        # safer alternative to numpy.asarray().
        #
        # TODO: move the dtype-casting to the view_converter's output space,
        #       once dtypes-for-spaces is merged into master.
        X = theano._asarray(X, theano.config.floatX)

        # Formats data as rows in a matrix, for DenseDesignMatrix
        X = X.reshape(-1, 2*numpy.prod(self.original_image_shape))

        # This is uint8
        y = SmallNORB.load(which_set, 'cat')
        if multi_target:
            y_extra = SmallNORB.load(which_set, 'info')
            y = numpy.hstack((y[:, numpy.newaxis], y_extra))

        datum_shape = ((2, ) +  # two stereo images
                       self.original_image_shape +
                       (1, ))  # one color channel

        # 's' is the stereo channel: 0 (left) or 1 (right)
        axes = ('b', 's', 0, 1, 'c')
        view_converter = StereoViewConverter(datum_shape, axes)

        super(SmallNORB, self).__init__(X=X,
                                        y=y,
                                        view_converter=view_converter)

    @classmethod
    def load(cls, which_set, filetype):
        """
        Reads and returns a single file as a numpy array.
        """

        assert which_set in ['train', 'test']
        assert filetype in ['dat', 'cat', 'info']

        def getPath(which_set):
            dirname = os.path.join(os.getenv('PYLEARN2_DATA_PATH'),
                                   'norb_small/original')
            if which_set == 'train':
                instance_list = '46789'
            elif which_set == 'test':
                instance_list = '01235'

            filename = 'smallnorb-5x%sx9x18x6x2x96x96-%s-%s.mat' % \
                (instance_list, which_set + 'ing', filetype)

            return os.path.join(dirname, filename)

        def parseNORBFile(file_handle, subtensor=None, debug=False):
            """
            Load all or part of file 'file_handle' into a numpy ndarray

            :param file_handle: file from which to read file can be opended
              with open(), gzip.open() and bz2.BZ2File()
              @type file_handle: file-like object. Can be a gzip open file.

            :param subtensor: If subtensor is not None, it should be like the
              argument to numpy.ndarray.__getitem__.  The following two
              expressions should return equivalent ndarray objects, but the one
              on the left may be faster and more memory efficient if the
              underlying file f is big.

              read(file_handle, subtensor) <===> read(file_handle)[*subtensor]

              Support for subtensors is currently spotty, so check the code to
              see if your particular type of subtensor is supported.
              """

            def readNums(file_handle, num_type, count):
                """
                Reads 4 bytes from file, returns it as a 32-bit integer.
                """
                num_bytes = count * numpy.dtype(num_type).itemsize
                string = file_handle.read(num_bytes)
                return numpy.fromstring(string, dtype=num_type)

            def readHeader(file_handle, debug=False, from_gzip=None):
                """
                :param file_handle: an open file handle.
                :type file_handle: a file or gzip.GzipFile object

                :param from_gzip: bool or None
                :type from_gzip: if None determine the type of file handle.

                :returns: data type, element size, rank, shape, size
                """

                if from_gzip is None:
                    from_gzip = isinstance(file_handle,
                                          (gzip.GzipFile, bz2.BZ2File))

                key_to_type = {0x1E3D4C51: ('float32', 4),
                               # what is a packed matrix?
                               # 0x1E3D4C52: ('packed matrix', 0),
                               0x1E3D4C53: ('float64', 8),
                               0x1E3D4C54: ('int32', 4),
                               0x1E3D4C55: ('uint8', 1),
                               0x1E3D4C56: ('int16', 2)}

                type_key = readNums(file_handle, 'int32', 1)[0]
                elem_type, elem_size = key_to_type[type_key]
                if debug:
                    print "header's type key, type, type size: ", \
                        type_key, elem_type, elem_size
                if elem_type == 'packed matrix':
                    raise NotImplementedError('packed matrix not supported')

                num_dims = readNums(file_handle, 'int32', 1)[0]
                if debug:
                    print '# of dimensions, according to header: ', num_dims

                if from_gzip:
                    shape = readNums(file_handle,
                                     'int32',
                                     max(num_dims, 3))[:num_dims]
                else:
                    shape = numpy.fromfile(file_handle,
                                           dtype='int32',
                                           count=max(num_dims, 3))[:num_dims]

                if debug:
                    print 'Tensor shape, as listed in header:', shape

                return elem_type, elem_size, shape

            elem_type, elem_size, shape = readHeader(file_handle, debug)
            beginning = file_handle.tell()

            num_elems = numpy.prod(shape)

            result = None
            if isinstance(file_handle, (gzip.GzipFile, bz2.BZ2File)):
                assert subtensor is None, \
                    "Subtensors on gzip files are not implemented."
                result = readNums(file_handle,
                                  elem_type,
                                  num_elems*elem_size).reshape(shape)
            elif subtensor is None:
                result = numpy.fromfile(file_handle,
                                        dtype=elem_type,
                                        count=num_elems).reshape(shape)
            elif isinstance(subtensor, slice):
                if subtensor.step not in (None, 1):
                    raise NotImplementedError('slice with step',
                                              subtensor.step)
                if subtensor.start not in (None, 0):
                    bytes_per_row = numpy.prod(shape[1:]) * elem_size
                    file_handle.seek(beginning+subtensor.start * bytes_per_row)
                shape[0] = min(shape[0], subtensor.stop) - subtensor.start
                result = numpy.fromfile(file_handle,
                                        dtype=elem_type,
                                        count=num_elems).reshape(shape)
            else:
                raise NotImplementedError('subtensor access not written yet:',
                                          subtensor)

            return result

        file_handle = open(getPath(which_set))
        return parseNORBFile(file_handle)

    def get_topological_view(self, mat=None, single_tensor=True):
        result = super(SmallNORB, self).get_topological_view(mat)

        if single_tensor:
            warnings.warn("The single_tensor argument is True by default to "
                          "maintain backwards compatibility. This argument "
                          "will be removed, and the behavior will become that "
                          "of single_tensor=False, as of August 2014.")
            axes = list(self.view_converter.axes)
            s_index = axes.index('s')
            assert axes.index('b') == 0
            num_image_pairs = result[0].shape[0]
            shape = (num_image_pairs, ) + self.view_converter.shape

            # inserts a singleton dimension where the 's' dimesion will be
            mono_shape = shape[:s_index] + (1, ) + shape[(s_index+1):]

            for i, res in enumerate(result):
                print "result %d shape: %s" % (i, str(res.shape))

            result = tuple(t.reshape(mono_shape) for t in result)
            result = numpy.concatenate(result, axis=s_index)
        else:
            warnings.warn("The single_tensor argument will be removed on "
                          "August 2014. The behavior will be the same as "
                          "single_tensor=False.")

        return result

class StereoViewConverter(object):
    """
    Converts stereo image data between two formats:
      A) A dense design matrix, one stereo pair per row (VectorSpace)
      B) An image pair (CompositeSpace of two Conv2DSpaces)
    """
    def __init__(self, shape, axes=None):
        """
        The arguments describe how the data is laid out in the design matrix.

        shape: A tuple of 4 ints, describing the shape of each datum.
               This is the size of each axis in <axes>, excluding the 'b' axis.

        axes: tuple of the following elements in any order:
          'b'  batch axis)
          's'  stereo axis)
           0   image axis 0 (row)
           1   image axis 1 (column)
          'c'  channel axis
        """
        shape = tuple(shape)

        if not all(isinstance(s, int) for s in shape):
            raise TypeError("Shape must be a tuple/list of ints")

        if len(shape) != 4:
            raise ValueError("Shape array needs to be of length 4, got %s." %
                             shape)

        datum_axes = list(axes)
        datum_axes.remove('b')
        if shape[datum_axes.index('s')] != 2:
            raise ValueError("Expected 's' axis to have size 2, got %d.\n"
                             "  axes:       %s\n"
                             "  shape:      %s" %
                             (shape[datum_axes.index('s')],
                              axes,
                              shape))
        self.shape = shape
        self.set_axes(axes)

        def make_conv2d_space(shape, axes):
            shape_axes = list(axes)
            shape_axes.remove('b')
            image_shape = tuple(shape[shape_axes.index(axis)]
                                for axis in (0, 1))
            conv2d_axes = list(axes)
            conv2d_axes.remove('s')
            return Conv2DSpace(shape=image_shape,
                               num_channels=shape[shape_axes.index('c')],
                               axes=conv2d_axes)

        conv2d_space = make_conv2d_space(shape, axes)
        self.topo_space = CompositeSpace((conv2d_space, conv2d_space))
        self.storage_space = VectorSpace(dim=numpy.prod(shape))

    def get_formatted_batch(self, batch, space):
        return self.storage_space.np_format_as(batch, space)

    def design_mat_to_topo_view(self, design_mat):
        """
        Called by DenseDesignMatrix.get_formatted_view(), get_batch_topo()
        """
        return self.storage_space.np_format_as(design_mat, self.topo_space)

    def design_mat_to_weights_view(self, design_mat):
        """
        Called by DenseDesignMatrix.get_weights_view()
        """
        return self.design_mat_to_topo_view(design_mat)

    def topo_view_to_design_mat(self, topo_batch):
        """
        Used by DenseDesignMatrix.set_topological_view(), .get_design_mat()
        """
        return self.topo_space.np_format_as(topo_batch, self.storage_space)

    def view_shape(self):
        return self.shape

    def weights_view_shape(self):
        return self.view_shape()

    def set_axes(self, axes):
        axes = tuple(axes)

        if len(axes) != 5:
            raise ValueError("Axes must have 5 elements; got %s" % str(axes))

        for required_axis in ('b', 's', 0, 1, 'c'):
            if required_axis not in axes:
                raise ValueError("Axes must contain 'b', 's', 0, 1, and 'c'. "
                                 "Got %s." % str(axes))

        if axes.index('b') != 0:
            raise ValueError("The 'b' axis must come first (axes = %s)." %
                             str(axes))

        def get_batchless_axes(axes):
            axes = list(axes)
            axes.remove('b')
            return tuple(axes)

        if hasattr(self, 'axes'):
            # Reorders the shape vector to match the new axis ordering.
            assert hasattr(self, 'shape')
            old_axes = get_batchless_axes(self.axes)
            new_axes = get_batchless_axes(axes)
            new_shape = tuple(self.shape[old_axes.index(a)] for a in new_axes)
            self.shape = new_shape

        self.axes = axes
<commit_msg>docstring correction<commit_after>"""
An interface to the small NORB dataset. Unlike ./norb_small.py, this reads the
original NORB file format, not the LISA lab's .npy version.

Currently only supports the Small NORB Dataset.

Download the dataset from:
http://www.cs.nyu.edu/~ylclab/data/norb-v1.0-small/

NORB dataset(s) by Fu Jie Huang and Yann LeCun.
"""

__authors__ = "Guillaume Desjardins and Matthew Koichi Grimes"
__copyright__ = "Copyright 2010-2014, Universite de Montreal"
__credits__ = __authors__.split(" and ")
__license__ = "3-clause BSD"
__maintainer__ = "Matthew Koichi Grimes"
__email__ = "mkg alum mit edu (@..)"

import os, gzip, bz2, warnings
import numpy, theano
from pylearn2.datasets import dense_design_matrix
from pylearn2.space import VectorSpace, Conv2DSpace, CompositeSpace

class SmallNORB(dense_design_matrix.DenseDesignMatrix):
    """
    An interface to the small NORB dataset.

    If instantiated with default arguments, target labels are integers
    representing categories, which can be looked up using

      category_name = SmallNORB.get_category(label).

    If instantiated with multi_target=True, labels are vectors of indices
    representing:

      [ category, instance, elevation, azimuth, lighting ]

    Like with category, there are class methods that map these ints to their
    actual values, e.g:

      category = SmallNORB.get_category(label[0])
      elevation = SmallNORB.get_elevation_degrees(label[2])
    """

    # Actual image shape may change, e.g. after being preprocessed by
    # datasets.preprocessing.Downsample
    original_image_shape = (96, 96)

    _categories = ['animal',  # four-legged animal
                   'human',  # human figure
                   'airplane',
                   'truck',
                   'car']

    @classmethod
    def get_category(cls, scalar_label):
        return cls._categories[int(scalar_label)]

    @classmethod
    def get_elevation_degrees(cls, scalar_label):
        scalar_label = int(scalar_label)
        assert scalar_label >= 0
        assert scalar_label < 9
        return 30 + 5 * scalar_label

    @classmethod
    def get_azimuth_degrees(cls, scalar_label):
        scalar_label = int(scalar_label)
        assert scalar_label >= 0
        assert scalar_label <= 34
        assert (scalar_label % 2) == 0
        return scalar_label * 10

    # Maps azimuth labels (ints) to their actual values, in degrees.
    azimuth_degrees = numpy.arange(0, 341, 20)

    # Maps a label type to its index within a label vector.
    label_type_to_index = {'category': 0,
                           'instance': 1,
                           'elevation': 2,
                           'azimuth': 3,
                           'lighting': 4}

    # Number of labels, for each label type.
    num_labels_by_type = (len(_categories),
                          10,  # instances
                          9,   # elevations
                          18,  # azimuths
                          6)   # lighting

    # [mkg] Dropped support for the 'center' argument for now. In Pylearn 1, it
    # shifted the pixel values from [0:255] by subtracting 127.5. Seems like a
    # form of preprocessing, which might be better implemented separately using
    # the Preprocess class.
    def __init__(self, which_set, multi_target=False):
        """
        parameters
        ----------

        which_set: str
            Must be 'train' or 'test'.

        multi_target: bool
            If False, each label is an integer labeling the image catergory. If
            True, each label is a vector: [category, instance, lighting,
            elevation, azimuth]. All labels are given as integers. Use the
            categories, elevation_degrees, and azimuth_degrees arrays to map
            from these integers to actual values.
        """

        assert which_set in ['train', 'test']

        self.which_set = which_set

        X = SmallNORB.load(which_set, 'dat')

        # Casts to the GPU-supported float type, using theano._asarray(), a
        # safer alternative to numpy.asarray().
        #
        # TODO: move the dtype-casting to the view_converter's output space,
        #       once dtypes-for-spaces is merged into master.
        X = theano._asarray(X, theano.config.floatX)

        # Formats data as rows in a matrix, for DenseDesignMatrix
        X = X.reshape(-1, 2*numpy.prod(self.original_image_shape))

        # This is uint8
        y = SmallNORB.load(which_set, 'cat')
        if multi_target:
            y_extra = SmallNORB.load(which_set, 'info')
            y = numpy.hstack((y[:, numpy.newaxis], y_extra))

        datum_shape = ((2, ) +  # two stereo images
                       self.original_image_shape +
                       (1, ))  # one color channel

        # 's' is the stereo channel: 0 (left) or 1 (right)
        axes = ('b', 's', 0, 1, 'c')
        view_converter = StereoViewConverter(datum_shape, axes)

        super(SmallNORB, self).__init__(X=X,
                                        y=y,
                                        view_converter=view_converter)

    @classmethod
    def load(cls, which_set, filetype):
        """
        Reads and returns a single file as a numpy array.
        """

        assert which_set in ['train', 'test']
        assert filetype in ['dat', 'cat', 'info']

        def getPath(which_set):
            dirname = os.path.join(os.getenv('PYLEARN2_DATA_PATH'),
                                   'norb_small/original')
            if which_set == 'train':
                instance_list = '46789'
            elif which_set == 'test':
                instance_list = '01235'

            filename = 'smallnorb-5x%sx9x18x6x2x96x96-%s-%s.mat' % \
                (instance_list, which_set + 'ing', filetype)

            return os.path.join(dirname, filename)

        def parseNORBFile(file_handle, subtensor=None, debug=False):
            """
            Load all or part of file 'file_handle' into a numpy ndarray

            :param file_handle: file from which to read file can be opended
              with open(), gzip.open() and bz2.BZ2File()
              @type file_handle: file-like object. Can be a gzip open file.

            :param subtensor: If subtensor is not None, it should be like the
              argument to numpy.ndarray.__getitem__.  The following two
              expressions should return equivalent ndarray objects, but the one
              on the left may be faster and more memory efficient if the
              underlying file f is big.

              read(file_handle, subtensor) <===> read(file_handle)[*subtensor]

              Support for subtensors is currently spotty, so check the code to
              see if your particular type of subtensor is supported.
              """

            def readNums(file_handle, num_type, count):
                """
                Reads 4 bytes from file, returns it as a 32-bit integer.
                """
                num_bytes = count * numpy.dtype(num_type).itemsize
                string = file_handle.read(num_bytes)
                return numpy.fromstring(string, dtype=num_type)

            def readHeader(file_handle, debug=False, from_gzip=None):
                """
                :param file_handle: an open file handle.
                :type file_handle: a file or gzip.GzipFile object

                :param from_gzip: bool or None
                :type from_gzip: if None determine the type of file handle.

                :returns: data type, element size, rank, shape, size
                """

                if from_gzip is None:
                    from_gzip = isinstance(file_handle,
                                          (gzip.GzipFile, bz2.BZ2File))

                key_to_type = {0x1E3D4C51: ('float32', 4),
                               # what is a packed matrix?
                               # 0x1E3D4C52: ('packed matrix', 0),
                               0x1E3D4C53: ('float64', 8),
                               0x1E3D4C54: ('int32', 4),
                               0x1E3D4C55: ('uint8', 1),
                               0x1E3D4C56: ('int16', 2)}

                type_key = readNums(file_handle, 'int32', 1)[0]
                elem_type, elem_size = key_to_type[type_key]
                if debug:
                    print "header's type key, type, type size: ", \
                        type_key, elem_type, elem_size
                if elem_type == 'packed matrix':
                    raise NotImplementedError('packed matrix not supported')

                num_dims = readNums(file_handle, 'int32', 1)[0]
                if debug:
                    print '# of dimensions, according to header: ', num_dims

                if from_gzip:
                    shape = readNums(file_handle,
                                     'int32',
                                     max(num_dims, 3))[:num_dims]
                else:
                    shape = numpy.fromfile(file_handle,
                                           dtype='int32',
                                           count=max(num_dims, 3))[:num_dims]

                if debug:
                    print 'Tensor shape, as listed in header:', shape

                return elem_type, elem_size, shape

            elem_type, elem_size, shape = readHeader(file_handle, debug)
            beginning = file_handle.tell()

            num_elems = numpy.prod(shape)

            result = None
            if isinstance(file_handle, (gzip.GzipFile, bz2.BZ2File)):
                assert subtensor is None, \
                    "Subtensors on gzip files are not implemented."
                result = readNums(file_handle,
                                  elem_type,
                                  num_elems*elem_size).reshape(shape)
            elif subtensor is None:
                result = numpy.fromfile(file_handle,
                                        dtype=elem_type,
                                        count=num_elems).reshape(shape)
            elif isinstance(subtensor, slice):
                if subtensor.step not in (None, 1):
                    raise NotImplementedError('slice with step',
                                              subtensor.step)
                if subtensor.start not in (None, 0):
                    bytes_per_row = numpy.prod(shape[1:]) * elem_size
                    file_handle.seek(beginning+subtensor.start * bytes_per_row)
                shape[0] = min(shape[0], subtensor.stop) - subtensor.start
                result = numpy.fromfile(file_handle,
                                        dtype=elem_type,
                                        count=num_elems).reshape(shape)
            else:
                raise NotImplementedError('subtensor access not written yet:',
                                          subtensor)

            return result

        file_handle = open(getPath(which_set))
        return parseNORBFile(file_handle)

    def get_topological_view(self, mat=None, single_tensor=True):
        result = super(SmallNORB, self).get_topological_view(mat)

        if single_tensor:
            warnings.warn("The single_tensor argument is True by default to "
                          "maintain backwards compatibility. This argument "
                          "will be removed, and the behavior will become that "
                          "of single_tensor=False, as of August 2014.")
            axes = list(self.view_converter.axes)
            s_index = axes.index('s')
            assert axes.index('b') == 0
            num_image_pairs = result[0].shape[0]
            shape = (num_image_pairs, ) + self.view_converter.shape

            # inserts a singleton dimension where the 's' dimesion will be
            mono_shape = shape[:s_index] + (1, ) + shape[(s_index+1):]

            for i, res in enumerate(result):
                print "result %d shape: %s" % (i, str(res.shape))

            result = tuple(t.reshape(mono_shape) for t in result)
            result = numpy.concatenate(result, axis=s_index)
        else:
            warnings.warn("The single_tensor argument will be removed on "
                          "August 2014. The behavior will be the same as "
                          "single_tensor=False.")

        return result

class StereoViewConverter(object):
    """
    Converts stereo image data between two formats:
      A) A dense design matrix, one stereo pair per row (VectorSpace)
      B) An image pair (CompositeSpace of two Conv2DSpaces)
    """
    def __init__(self, shape, axes=None):
        """
        The arguments describe how the data is laid out in the design matrix.

        shape: A tuple of 4 ints, describing the shape of each datum.
               This is the size of each axis in <axes>, excluding the 'b' axis.

        axes: tuple of the following elements in any order:
          'b'  batch axis)
          's'  stereo axis)
           0   image axis 0 (row)
           1   image axis 1 (column)
          'c'  channel axis
        """
        shape = tuple(shape)

        if not all(isinstance(s, int) for s in shape):
            raise TypeError("Shape must be a tuple/list of ints")

        if len(shape) != 4:
            raise ValueError("Shape array needs to be of length 4, got %s." %
                             shape)

        datum_axes = list(axes)
        datum_axes.remove('b')
        if shape[datum_axes.index('s')] != 2:
            raise ValueError("Expected 's' axis to have size 2, got %d.\n"
                             "  axes:       %s\n"
                             "  shape:      %s" %
                             (shape[datum_axes.index('s')],
                              axes,
                              shape))
        self.shape = shape
        self.set_axes(axes)

        def make_conv2d_space(shape, axes):
            shape_axes = list(axes)
            shape_axes.remove('b')
            image_shape = tuple(shape[shape_axes.index(axis)]
                                for axis in (0, 1))
            conv2d_axes = list(axes)
            conv2d_axes.remove('s')
            return Conv2DSpace(shape=image_shape,
                               num_channels=shape[shape_axes.index('c')],
                               axes=conv2d_axes)

        conv2d_space = make_conv2d_space(shape, axes)
        self.topo_space = CompositeSpace((conv2d_space, conv2d_space))
        self.storage_space = VectorSpace(dim=numpy.prod(shape))

    def get_formatted_batch(self, batch, space):
        return self.storage_space.np_format_as(batch, space)

    def design_mat_to_topo_view(self, design_mat):
        """
        Called by DenseDesignMatrix.get_formatted_view(), get_batch_topo()
        """
        return self.storage_space.np_format_as(design_mat, self.topo_space)

    def design_mat_to_weights_view(self, design_mat):
        """
        Called by DenseDesignMatrix.get_weights_view()
        """
        return self.design_mat_to_topo_view(design_mat)

    def topo_view_to_design_mat(self, topo_batch):
        """
        Used by DenseDesignMatrix.set_topological_view(), .get_design_mat()
        """
        return self.topo_space.np_format_as(topo_batch, self.storage_space)

    def view_shape(self):
        return self.shape

    def weights_view_shape(self):
        return self.view_shape()

    def set_axes(self, axes):
        axes = tuple(axes)

        if len(axes) != 5:
            raise ValueError("Axes must have 5 elements; got %s" % str(axes))

        for required_axis in ('b', 's', 0, 1, 'c'):
            if required_axis not in axes:
                raise ValueError("Axes must contain 'b', 's', 0, 1, and 'c'. "
                                 "Got %s." % str(axes))

        if axes.index('b') != 0:
            raise ValueError("The 'b' axis must come first (axes = %s)." %
                             str(axes))

        def get_batchless_axes(axes):
            axes = list(axes)
            axes.remove('b')
            return tuple(axes)

        if hasattr(self, 'axes'):
            # Reorders the shape vector to match the new axis ordering.
            assert hasattr(self, 'shape')
            old_axes = get_batchless_axes(self.axes)
            new_axes = get_batchless_axes(axes)
            new_shape = tuple(self.shape[old_axes.index(a)] for a in new_axes)
            self.shape = new_shape

        self.axes = axes
<|endoftext|>