<commit_before># Copyright (c) 2016-2017 ARM Limited
# All rights reserved.
#
# The license below extends only to copyright in the software and shall
# not be construed as granting a license to any other intellectual
# property including but not limited to intellectual property relating
# to a hardware implementation of the functionality of the software
# licensed hereunder.  You may use the software subject to the license
# terms below provided that you ensure that this notice is replicated
# unmodified and in its entirety in all distributions of the software,
# modified or unmodified, in source code or in binary form.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met: redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer;
# redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution;
# neither the name of the copyright holders nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Andreas Sandberg
#          Gabor Dozsa

# System components used by the bigLITTLE.py configuration script

import m5
from m5.objects import *
m5.util.addToPath('../../')
from common.Caches import *
from common import CpuConfig

have_kvm = "ArmV8KvmCPU" in CpuConfig.cpu_names()

class L1I(L1_ICache):
    tag_latency = 1
    data_latency = 1
    response_latency = 1
    mshrs = 4
    tgts_per_mshr = 8
    size = '48kB'
    assoc = 3

class L1D(L1_DCache):
    tag_latency = 2
    data_latency = 2
    response_latency = 1
    mshrs = 16
    tgts_per_mshr = 16
    size = '32kB'
    assoc = 2
    write_buffers = 16

class WalkCache(PageTableWalkerCache):
    tag_latency = 4
    data_latency = 4
    response_latency = 4
    mshrs = 6
    tgts_per_mshr = 8
    size = '1kB'
    assoc = 8
    write_buffers = 16

class L2(L2Cache):
    tag_latency = 12
    data_latency = 12
    response_latency = 5
    mshrs = 32
    tgts_per_mshr = 8
    size = '1MB'
    assoc = 16
    write_buffers = 8
    clusivity='mostly_excl'

class L3(Cache):
    size = '16MB'
    assoc = 16
    tag_latency = 20
    data_latency = 20
    response_latency = 20
    mshrs = 20
    tgts_per_mshr = 12
    clusivity='mostly_excl'

class MemBus(SystemXBar):
    badaddr_responder = BadAddr(warn_access="warn")
    default = Self.badaddr_responder.pio

class CpuCluster(SubSystem):
    def __init__(self, system,  num_cpus, cpu_clock, cpu_voltage,
                 cpu_type, l1i_type, l1d_type, wcache_type, l2_type):
        super(CpuCluster, self).__init__()
        self._cpu_type = cpu_type
        self._l1i_type = l1i_type
        self._l1d_type = l1d_type
        self._wcache_type = wcache_type
        self._l2_type = l2_type

        assert num_cpus > 0

        self.voltage_domain = VoltageDomain(voltage=cpu_voltage)
        self.clk_domain = SrcClockDomain(clock=cpu_clock,
                                         voltage_domain=self.voltage_domain)

        self.cpus = [ self._cpu_type(cpu_id=system.numCpus() + idx,
                                     clk_domain=self.clk_domain)
                      for idx in range(num_cpus) ]

        for cpu in self.cpus:
            cpu.createThreads()
            cpu.createInterruptController()
            cpu.socket_id = system.numCpuClusters()
        system.addCpuCluster(self, num_cpus)

    def requireCaches(self):
        return self._cpu_type.require_caches()

    def memoryMode(self):
        return self._cpu_type.memory_mode()

    def addL1(self):
        for cpu in self.cpus:
            l1i = None if self._l1i_type is None else self._l1i_type()
            l1d = None if self._l1d_type is None else self._l1d_type()
            iwc = None if self._wcache_type is None else self._wcache_type()
            dwc = None if self._wcache_type is None else self._wcache_type()
            cpu.addPrivateSplitL1Caches(l1i, l1d, iwc, dwc)

    def addL2(self, clk_domain):
        if self._l2_type is None:
            return
        self.toL2Bus = L2XBar(width=64, clk_domain=clk_domain)
        self.l2 = self._l2_type()
        for cpu in self.cpus:
            cpu.connectAllPorts(self.toL2Bus)
        self.toL2Bus.master = self.l2.cpu_side

    def connectMemSide(self, bus):
        bus.slave
        try:
            self.l2.mem_side = bus.slave
        except AttributeError:
            for cpu in self.cpus:
                cpu.connectAllPorts(bus)

class AtomicCluster(CpuCluster):
    def __init__(self, system, num_cpus, cpu_clock, cpu_voltage="1.0V"):
        cpu_config = [ CpuConfig.get("AtomicSimpleCPU"), None, None, None, None ]
        super(AtomicCluster, self).__init__(system, num_cpus, cpu_clock,
                                            cpu_voltage, *cpu_config)
    def addL1(self):
        pass

class KvmCluster(CpuCluster):
    def __init__(self, system, num_cpus, cpu_clock, cpu_voltage="1.0V"):
        cpu_config = [ CpuConfig.get("ArmV8KvmCPU"), None, None, None, None ]
        super(KvmCluster, self).__init__(system, num_cpus, cpu_clock,
                                         cpu_voltage, *cpu_config)
    def addL1(self):
        pass

class SimpleSystem(LinuxArmSystem):
    cache_line_size = 64

    def __init__(self, caches, mem_size, **kwargs):
        super(SimpleSystem, self).__init__(**kwargs)

        self.voltage_domain = VoltageDomain(voltage="1.0V")
        self.clk_domain = SrcClockDomain(clock="1GHz",
                                         voltage_domain=Parent.voltage_domain)

        self.realview = VExpress_GEM5_V1()

        self.gic_cpu_addr = self.realview.gic.cpu_addr
        self.flags_addr = self.realview.realview_io.pio_addr + 0x30

        self.membus = MemBus()

        self.intrctrl = IntrControl()
        self.terminal = Terminal()
        self.vncserver = VncServer()

        self.iobus = IOXBar()
        # CPUs->PIO
        self.iobridge = Bridge(delay='50ns')
        # Device DMA -> MEM
        mem_range = self.realview._mem_regions[0]
        mem_range_size = long(mem_range[1]) - long(mem_range[0])
        assert mem_range_size >= long(Addr(mem_size))
        self.mem_ranges = [ AddrRange(start=mem_range[0], size=mem_size) ]
        self._caches = caches
        if self._caches:
            self.iocache = IOCache(addr_ranges=[self.mem_ranges[0]])
        else:
            self.dmabridge = Bridge(delay='50ns',
                                    ranges=[self.mem_ranges[0]])

        self._pci_devices = 0
        self._clusters = []
        self._num_cpus = 0

    def attach_pci(self, dev):
        dev.pci_bus, dev.pci_dev, dev.pci_func = (0, self._pci_devices + 1, 0)
        self._pci_devices += 1
        self.realview.attachPciDevice(dev, self.iobus)

    def connect(self):
        self.iobridge.master = self.iobus.slave
        self.iobridge.slave = self.membus.master

        if self._caches:
            self.iocache.mem_side = self.membus.slave
            self.iocache.cpu_side = self.iobus.master
        else:
            self.dmabridge.master = self.membus.slave
            self.dmabridge.slave = self.iobus.master

        self.gic_cpu_addr = self.realview.gic.cpu_addr
        self.realview.attachOnChipIO(self.membus, self.iobridge)
        self.realview.attachIO(self.iobus)
        self.system_port = self.membus.slave

    def numCpuClusters(self):
        return len(self._clusters)

    def addCpuCluster(self, cpu_cluster, num_cpus):
        assert cpu_cluster not in self._clusters
        assert num_cpus > 0
        self._clusters.append(cpu_cluster)
        self._num_cpus += num_cpus

    def numCpus(self):
        return self._num_cpus

    def addCaches(self, need_caches, last_cache_level):
        if not need_caches:
            # connect each cluster to the memory hierarchy
            for cluster in self._clusters:
                cluster.connectMemSide(self.membus)
            return

        cluster_mem_bus = self.membus
        assert last_cache_level >= 1 and last_cache_level <= 3
        for cluster in self._clusters:
            cluster.addL1()
        if last_cache_level > 1:
            for cluster in self._clusters:
                cluster.addL2(cluster.clk_domain)
        if last_cache_level > 2:
            max_clock_cluster = max(self._clusters,
                                    key=lambda c: c.clk_domain.clock[0])
            self.l3 = L3(clk_domain=max_clock_cluster.clk_domain)
            self.toL3Bus = L2XBar(width=64)
            self.toL3Bus.master = self.l3.cpu_side
            self.l3.mem_side = self.membus.slave
            cluster_mem_bus = self.toL3Bus

        # connect each cluster to the memory hierarchy
        for cluster in self._clusters:
            cluster.connectMemSide(cluster_mem_bus)
<commit_msg>configs: Unifiy interpretation of Realview mem_regions<commit_after># Copyright (c) 2016-2017 ARM Limited
# All rights reserved.
#
# The license below extends only to copyright in the software and shall
# not be construed as granting a license to any other intellectual
# property including but not limited to intellectual property relating
# to a hardware implementation of the functionality of the software
# licensed hereunder.  You may use the software subject to the license
# terms below provided that you ensure that this notice is replicated
# unmodified and in its entirety in all distributions of the software,
# modified or unmodified, in source code or in binary form.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met: redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer;
# redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution;
# neither the name of the copyright holders nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Andreas Sandberg
#          Gabor Dozsa

# System components used by the bigLITTLE.py configuration script

import m5
from m5.objects import *
m5.util.addToPath('../../')
from common.Caches import *
from common import CpuConfig

have_kvm = "ArmV8KvmCPU" in CpuConfig.cpu_names()

class L1I(L1_ICache):
    tag_latency = 1
    data_latency = 1
    response_latency = 1
    mshrs = 4
    tgts_per_mshr = 8
    size = '48kB'
    assoc = 3

class L1D(L1_DCache):
    tag_latency = 2
    data_latency = 2
    response_latency = 1
    mshrs = 16
    tgts_per_mshr = 16
    size = '32kB'
    assoc = 2
    write_buffers = 16

class WalkCache(PageTableWalkerCache):
    tag_latency = 4
    data_latency = 4
    response_latency = 4
    mshrs = 6
    tgts_per_mshr = 8
    size = '1kB'
    assoc = 8
    write_buffers = 16

class L2(L2Cache):
    tag_latency = 12
    data_latency = 12
    response_latency = 5
    mshrs = 32
    tgts_per_mshr = 8
    size = '1MB'
    assoc = 16
    write_buffers = 8
    clusivity='mostly_excl'

class L3(Cache):
    size = '16MB'
    assoc = 16
    tag_latency = 20
    data_latency = 20
    response_latency = 20
    mshrs = 20
    tgts_per_mshr = 12
    clusivity='mostly_excl'

class MemBus(SystemXBar):
    badaddr_responder = BadAddr(warn_access="warn")
    default = Self.badaddr_responder.pio

class CpuCluster(SubSystem):
    def __init__(self, system,  num_cpus, cpu_clock, cpu_voltage,
                 cpu_type, l1i_type, l1d_type, wcache_type, l2_type):
        super(CpuCluster, self).__init__()
        self._cpu_type = cpu_type
        self._l1i_type = l1i_type
        self._l1d_type = l1d_type
        self._wcache_type = wcache_type
        self._l2_type = l2_type

        assert num_cpus > 0

        self.voltage_domain = VoltageDomain(voltage=cpu_voltage)
        self.clk_domain = SrcClockDomain(clock=cpu_clock,
                                         voltage_domain=self.voltage_domain)

        self.cpus = [ self._cpu_type(cpu_id=system.numCpus() + idx,
                                     clk_domain=self.clk_domain)
                      for idx in range(num_cpus) ]

        for cpu in self.cpus:
            cpu.createThreads()
            cpu.createInterruptController()
            cpu.socket_id = system.numCpuClusters()
        system.addCpuCluster(self, num_cpus)

    def requireCaches(self):
        return self._cpu_type.require_caches()

    def memoryMode(self):
        return self._cpu_type.memory_mode()

    def addL1(self):
        for cpu in self.cpus:
            l1i = None if self._l1i_type is None else self._l1i_type()
            l1d = None if self._l1d_type is None else self._l1d_type()
            iwc = None if self._wcache_type is None else self._wcache_type()
            dwc = None if self._wcache_type is None else self._wcache_type()
            cpu.addPrivateSplitL1Caches(l1i, l1d, iwc, dwc)

    def addL2(self, clk_domain):
        if self._l2_type is None:
            return
        self.toL2Bus = L2XBar(width=64, clk_domain=clk_domain)
        self.l2 = self._l2_type()
        for cpu in self.cpus:
            cpu.connectAllPorts(self.toL2Bus)
        self.toL2Bus.master = self.l2.cpu_side

    def connectMemSide(self, bus):
        bus.slave
        try:
            self.l2.mem_side = bus.slave
        except AttributeError:
            for cpu in self.cpus:
                cpu.connectAllPorts(bus)

class AtomicCluster(CpuCluster):
    def __init__(self, system, num_cpus, cpu_clock, cpu_voltage="1.0V"):
        cpu_config = [ CpuConfig.get("AtomicSimpleCPU"), None, None, None, None ]
        super(AtomicCluster, self).__init__(system, num_cpus, cpu_clock,
                                            cpu_voltage, *cpu_config)
    def addL1(self):
        pass

class KvmCluster(CpuCluster):
    def __init__(self, system, num_cpus, cpu_clock, cpu_voltage="1.0V"):
        cpu_config = [ CpuConfig.get("ArmV8KvmCPU"), None, None, None, None ]
        super(KvmCluster, self).__init__(system, num_cpus, cpu_clock,
                                         cpu_voltage, *cpu_config)
    def addL1(self):
        pass

class SimpleSystem(LinuxArmSystem):
    cache_line_size = 64

    def __init__(self, caches, mem_size, **kwargs):
        super(SimpleSystem, self).__init__(**kwargs)

        self.voltage_domain = VoltageDomain(voltage="1.0V")
        self.clk_domain = SrcClockDomain(clock="1GHz",
                                         voltage_domain=Parent.voltage_domain)

        self.realview = VExpress_GEM5_V1()

        self.gic_cpu_addr = self.realview.gic.cpu_addr
        self.flags_addr = self.realview.realview_io.pio_addr + 0x30

        self.membus = MemBus()

        self.intrctrl = IntrControl()
        self.terminal = Terminal()
        self.vncserver = VncServer()

        self.iobus = IOXBar()
        # CPUs->PIO
        self.iobridge = Bridge(delay='50ns')
        # Device DMA -> MEM
        mem_range = self.realview._mem_regions[0]
        max_size = long(mem_range[1])
        assert max_size >= long(Addr(mem_size))
        self.mem_ranges = [ AddrRange(start=mem_range[0], size=mem_size) ]
        self._caches = caches
        if self._caches:
            self.iocache = IOCache(addr_ranges=[self.mem_ranges[0]])
        else:
            self.dmabridge = Bridge(delay='50ns',
                                    ranges=[self.mem_ranges[0]])

        self._pci_devices = 0
        self._clusters = []
        self._num_cpus = 0

    def attach_pci(self, dev):
        dev.pci_bus, dev.pci_dev, dev.pci_func = (0, self._pci_devices + 1, 0)
        self._pci_devices += 1
        self.realview.attachPciDevice(dev, self.iobus)

    def connect(self):
        self.iobridge.master = self.iobus.slave
        self.iobridge.slave = self.membus.master

        if self._caches:
            self.iocache.mem_side = self.membus.slave
            self.iocache.cpu_side = self.iobus.master
        else:
            self.dmabridge.master = self.membus.slave
            self.dmabridge.slave = self.iobus.master

        self.gic_cpu_addr = self.realview.gic.cpu_addr
        self.realview.attachOnChipIO(self.membus, self.iobridge)
        self.realview.attachIO(self.iobus)
        self.system_port = self.membus.slave

    def numCpuClusters(self):
        return len(self._clusters)

    def addCpuCluster(self, cpu_cluster, num_cpus):
        assert cpu_cluster not in self._clusters
        assert num_cpus > 0
        self._clusters.append(cpu_cluster)
        self._num_cpus += num_cpus

    def numCpus(self):
        return self._num_cpus

    def addCaches(self, need_caches, last_cache_level):
        if not need_caches:
            # connect each cluster to the memory hierarchy
            for cluster in self._clusters:
                cluster.connectMemSide(self.membus)
            return

        cluster_mem_bus = self.membus
        assert last_cache_level >= 1 and last_cache_level <= 3
        for cluster in self._clusters:
            cluster.addL1()
        if last_cache_level > 1:
            for cluster in self._clusters:
                cluster.addL2(cluster.clk_domain)
        if last_cache_level > 2:
            max_clock_cluster = max(self._clusters,
                                    key=lambda c: c.clk_domain.clock[0])
            self.l3 = L3(clk_domain=max_clock_cluster.clk_domain)
            self.toL3Bus = L2XBar(width=64)
            self.toL3Bus.master = self.l3.cpu_side
            self.l3.mem_side = self.membus.slave
            cluster_mem_bus = self.toL3Bus

        # connect each cluster to the memory hierarchy
        for cluster in self._clusters:
            cluster.connectMemSide(cluster_mem_bus)
<|endoftext|>


<commit_before>'''
Show or edit MicroDrop configuration.

.. versionadded:: 2.13
'''
import argparse
import io
import json
import re
import sys

import configobj
import microdrop as md
import microdrop.config
import pydash
import yaml

def _config_parser():
    parser = argparse.ArgumentParser(parents=[md.MICRODROP_PARSER],
                                     add_help=False)

    subparsers = parser.add_subparsers(dest='command', help='commands')

    subparsers.add_parser('locate', help='Show path to configuration '
                          'source')

    show = subparsers.add_parser('show', help='Show configuration')
    show.add_argument('--get', metavar='KEY')
    show_format = show.add_mutually_exclusive_group()
    show_format.add_argument('--json', action='store_true')
    show_format.add_argument('--yaml', action='store_true')

    edit = subparsers.add_parser('edit', help='Modify configuration')
    edit.add_argument('-n', '--dry-run', action='store_true')
    edit_args = edit.add_mutually_exclusive_group(required=True)
    edit_args.add_argument('--append', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Add one configuration value to the beginning '
                           'of a list key.')
    edit_args.add_argument('--prepend', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Add one configuration value to the end of a '
                           'list key.')
    edit_args.add_argument('--set', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Set a boolean or string key')
    edit_args.add_argument('--remove', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Remove all instances of configuration value '
                           'from a list key.')
    edit_args.add_argument('--remove-key', metavar='KEY', help='Remove a '
                           'configuration key (and all its values).')
    return parser

CONFIG_PARSER = _config_parser()

def parse_args(args=None):
    """Parses arguments, returns (options, args)."""
    if args is None:
        args = sys.argv[1:]

    parser = argparse.ArgumentParser(parents=[CONFIG_PARSER])
    args = parser.parse_args(args)
    return args

def main(args=None):
    if args is None:
        args = parse_args()

    config = md.config.Config(args.config)

    if args.command == 'locate':
        print config.filename
    elif args.command == 'show':
        if args.get:
            data = pydash.get(config.data.dict(), args.get)
        else:
            data = config.data.dict()

        if args.json:
            # Output in JSON.
            json.dump(obj=data, fp=sys.stdout, indent=4)
        elif args.yaml:
            # Output in YAML format.
            print yaml.dump(data, default_flow_style=False),
        elif isinstance(data, dict):
            # Output in `ini` format.
            output = io.BytesIO()
            configobj.ConfigObj(data).write(output)
            print output.getvalue(),
        else:
            print data
    elif args.command == 'edit':
        for action_i in ('append', 'prepend', 'set', 'remove', 'remove_key'):
            if getattr(args, action_i):
                action = action_i
                break

        if action in ('append', 'prepend', 'set', 'remove'):
            # Unpack key and new value.
            key, new_value = getattr(args, action)

            # Look up existing value.
            config_value = pydash.get(config.data, key)

            if action in ('append', 'prepend', 'remove'):
                # Action is a list action.

                if config_value is None:
                    # Create dictionary structure containing only empty list for
                    # specified key.
                    config_value = []
                    nested_value = pydash.set_({}, key, config_value)
                    # Merge nested value into existing configuration structure.
                    pydash.merge(config.data, nested_value)
                elif not isinstance(config_value, list):
                    print >> sys.stderr, 'Value at %s is not a list.' % key
                    raise SystemExit(1)

                if new_value in config_value:
                    # Remove value even if we are appending or prepending to
                    # avoid duplicate values.
                    config_value.remove(new_value)

                if args.append:
                    config_value.append(new_value)
                elif args.prepend:
                    config_value.insert(0, new_value)
        elif action == 'remove_key':
            key = getattr(args, action)

            if pydash.get(config.data, key) is not None:
                # Key exists.

                # Split key into levels.
                # Use [negative lookbehind assertion][1] to only split on
                # non-escaped '.' characters.
                #
                # [1]: https://stackoverflow.com/a/21107911/345236
                levels = re.split(r'(?<!\\)\.', key)
                parents = levels[:-1]

                parent = config.data

                for parent_i in parents:
                    parent = parent[parent_i]

                # Delete key from deepest parent.
                del parent[levels[-1]]
        if args.dry_run:
            output = io.BytesIO()
            config.data.write(output)
            print output.getvalue(),
        else:
            config.save()
    return config

if __name__ == '__main__':
    config = main()
<commit_msg>[Fix] Use integer return code for bin.config:main<commit_after>'''
Show or edit MicroDrop configuration.

.. versionadded:: 2.13
'''
import argparse
import io
import json
import re
import sys

import configobj
import microdrop as md
import microdrop.config
import pydash
import yaml

def _config_parser():
    parser = argparse.ArgumentParser(parents=[md.MICRODROP_PARSER],
                                     add_help=False)

    subparsers = parser.add_subparsers(dest='command', help='commands')

    subparsers.add_parser('locate', help='Show path to configuration '
                          'source')

    show = subparsers.add_parser('show', help='Show configuration')
    show.add_argument('--get', metavar='KEY')
    show_format = show.add_mutually_exclusive_group()
    show_format.add_argument('--json', action='store_true')
    show_format.add_argument('--yaml', action='store_true')

    edit = subparsers.add_parser('edit', help='Modify configuration')
    edit.add_argument('-n', '--dry-run', action='store_true')
    edit_args = edit.add_mutually_exclusive_group(required=True)
    edit_args.add_argument('--append', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Add one configuration value to the beginning '
                           'of a list key.')
    edit_args.add_argument('--prepend', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Add one configuration value to the end of a '
                           'list key.')
    edit_args.add_argument('--set', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Set a boolean or string key')
    edit_args.add_argument('--remove', nargs=2, metavar=('KEY', 'VALUE'),
                           help='Remove all instances of configuration value '
                           'from a list key.')
    edit_args.add_argument('--remove-key', metavar='KEY', help='Remove a '
                           'configuration key (and all its values).')
    return parser

CONFIG_PARSER = _config_parser()

def parse_args(args=None):
    """Parses arguments, returns (options, args)."""
    if args is None:
        args = sys.argv[1:]

    parser = argparse.ArgumentParser(parents=[CONFIG_PARSER])
    args = parser.parse_args(args)
    return args

def main(args=None):
    '''
    Wrap :func:`config` with integer return code.

    Parameters
    ----------
    args : argparse.Namespace, optional
        Arguments as parsed by :func:`parse_args`.

    See also
    --------
    :func:`parse_args`
    '''
    config(args)
    # Return
    return 0

def config(args=None):
    '''
    Parameters
    ----------
    args : argparse.Namespace, optional
        Arguments as parsed by :func:`parse_args`.

    See also
    --------
    :func:`parse_args`

    Returns
    -------
    configobj.ConfigObj
        Parsed (and potentially modified) configuration.
    '''
    if args is None:
        args = parse_args()

    config = md.config.Config(args.config)

    if args.command == 'locate':
        print config.filename
    elif args.command == 'show':
        if args.get:
            data = pydash.get(config.data.dict(), args.get)
        else:
            data = config.data.dict()

        if args.json:
            # Output in JSON.
            json.dump(obj=data, fp=sys.stdout, indent=4)
        elif args.yaml:
            # Output in YAML format.
            print yaml.dump(data, default_flow_style=False),
        elif isinstance(data, dict):
            # Output in `ini` format.
            output = io.BytesIO()
            configobj.ConfigObj(data).write(output)
            print output.getvalue(),
        else:
            print data
    elif args.command == 'edit':
        for action_i in ('append', 'prepend', 'set', 'remove', 'remove_key'):
            if getattr(args, action_i):
                action = action_i
                break

        if action in ('append', 'prepend', 'set', 'remove'):
            # Unpack key and new value.
            key, new_value = getattr(args, action)

            # Look up existing value.
            config_value = pydash.get(config.data, key)

            if action in ('append', 'prepend', 'remove'):
                # Action is a list action.

                if config_value is None:
                    # Create dictionary structure containing only empty list for
                    # specified key.
                    config_value = []
                    nested_value = pydash.set_({}, key, config_value)
                    # Merge nested value into existing configuration structure.
                    pydash.merge(config.data, nested_value)
                elif not isinstance(config_value, list):
                    print >> sys.stderr, 'Value at %s is not a list.' % key
                    raise SystemExit(1)

                if new_value in config_value:
                    # Remove value even if we are appending or prepending to
                    # avoid duplicate values.
                    config_value.remove(new_value)

                if args.append:
                    config_value.append(new_value)
                elif args.prepend:
                    config_value.insert(0, new_value)
        elif action == 'remove_key':
            key = getattr(args, action)

            if pydash.get(config.data, key) is not None:
                # Key exists.

                # Split key into levels.
                # Use [negative lookbehind assertion][1] to only split on
                # non-escaped '.' characters.
                #
                # [1]: https://stackoverflow.com/a/21107911/345236
                levels = re.split(r'(?<!\\)\.', key)
                parents = levels[:-1]

                parent = config.data

                for parent_i in parents:
                    parent = parent[parent_i]

                # Delete key from deepest parent.
                del parent[levels[-1]]
        if args.dry_run:
            output = io.BytesIO()
            config.data.write(output)
            print output.getvalue(),
        else:
            config.save()
    return config

if __name__ == '__main__':
    config_ = config(parse_args())
<|endoftext|>


<commit_before>#!/usr/bin/python
""" Classes and functions for fitting ivim model """
from __future__ import division, print_function, absolute_import
import warnings
import functools
import numpy as np
from dipy.core.optimize import Optimizer

from scipy.optimize import leastsq

from dipy.core.gradients import gradient_table
from dipy.reconst.base import ReconstModel
from dipy.reconst.dti import _min_positive_signal
from dipy.reconst.dti import TensorModel, mean_diffusivity

def ivim_function(params, bvals):
    """The ivim function
    Parameters
    ----------
    params : array
             The IVIM parameters [S0, f, D_star, D]

    bvals  : array
             bvalues
    """
    S0, f, D_star, D = params
    S = S0 * (f * np.exp(-bvals * D_star) + (1 - f) * np.exp(-bvals * D))
    return S

def _ivim_error(params, bvals, signal):
    """Error function to be used in fitting the model
    """
    return (signal - ivim_function(params, bvals))**2

class IvimModel(ReconstModel):
    """Ivim model
    """

    def __init__(self, gtab, split_b=200.0, min_signal=None):
        """
        An IVIM model

        Parameters
        ----------
        gtab : GradientTable class instance
        split_b : split the data at this b-value for a two stage fit

        References
        ----------
        .. [1] Le Bihan, Denis, et al. "Separation of diffusion
                   and perfusion in intravoxel incoherent motion MR
                   imaging." Radiology 168.2 (1988): 497-505.
        .. [2] Federau, Christian, et al. "Quantitative measurement
                   of brain perfusion with intravoxel incoherent motion
                   MR imaging." Radiology 265.3 (2012): 874-881.
        """
        ReconstModel.__init__(self, gtab)
        self.split_b = split_b
        self.min_signal = min_signal
        if self.min_signal is not None and self.min_signal <= 0:
            e_s = "The `min_signal` key-word argument needs to be strictly"
            e_s += " positive."
            raise ValueError(e_s)

    def fit(self, data, mask=None, x0=None, fit_method="one_stage", routine="leastsq",
            jac=False, bounds=((0, 1.), (0, 1.), (0, 1.), (0, 1.)), tol=1e-25, algorithm='L-BFGS-B',
            gtol=1e-25, ftol=1e-25, eps=1e-15):
        """ Fit method of the Ivim model class

        Parameters
        ----------
        data : array
            The measured signal from one voxel.
        mask : array
            A boolean array used to mark the coordinates in the data that
            should be analyzed that has the shape data.shape[:-1]
        x0 : array
            Initial guesses for the parameters S0, f, D_star and D
            Default x0 = [1.0, 0.10, 0.001, 0.0009]
            These guess parameters are taken from the Federau paper
            Dimension can either be 1 or (N, 4) where N is the number of
            voxels in the data.
        fit_method: str
            Use one-stage fitting or two-stage fitting.
        jac : Boolean
            Use the Jacobian
        bounds : tuple
            Bounds for the various parameters
        tol : float
            tolerance for convergence
        method : str
            Fitting algorithm to use from scipy.optimize

        Returns
        -------
        IvimFit object
        """
        if mask is None:
            # Flatten it to 2D either way:
            data_in_mask = np.reshape(data, (-1, data.shape[-1]))
        else:
            # Check for valid shape of the mask
            if mask.shape != data.shape[:-1]:
                raise ValueError("Mask is not the same shape as data.")
            mask = np.array(mask, dtype=bool, copy=False)
            data_in_mask = np.reshape(data[mask], (-1, data.shape[-1]))

        if self.min_signal is None:
            min_signal = _min_positive_signal(data)
        else:
            min_signal = self.min_signal

        # Generate guess_parameters for all voxels
        if x0 is None:
            x0 = np.ones(
                (data.shape[:-1] + (4,))) * [1.0, 0.10, 0.001, 0.0009]
        else:
            if x0.shape != (data.shape[0], 4):
                raise ValueError(
                    "Guess params should be of shape (num_voxels,4)")
        data_in_mask = np.maximum(data_in_mask, min_signal)

        if fit_method == "one_stage":
            params_in_mask = one_stage(data_in_mask, self.gtab, x0,
                                       jac, bounds, tol,
                                       routine,
                                       algorithm, gtol, ftol, eps)
        elif fit_method == "two_stage":
            params_in_mask = two_stage(data_in_mask, self.gtab, x0,
                                       self.split_b, jac, bounds, tol,
                                       routine,
                                       algorithm, gtol, ftol, eps)
        else:
            raise ValueError("""Fit method must be either
                             'one_stage' or 'two_stage'""")

        if mask is None:
            out_shape = data.shape[:-1] + (-1, )
            ivim_params = params_in_mask.reshape(out_shape)
        else:
            ivim_params = np.zeros(data.shape[:-1] + (4,))
            ivim_params[mask, :] = params_in_mask

        return IvimFit(self, ivim_params)

class IvimFit(object):

    def __init__(self, model, model_params):
        """ Initialize a IvimFit class instance.
            Parameters
            ----------
            The model parameters are S0, f, D_star, D
        """
        self.model = model
        self.model_params = model_params

    def __getitem__(self, index):
        model_params = self.model_params
        N = model_params.ndim
        if type(index) is not tuple:
            index = (index,)
        elif len(index) >= model_params.ndim:
            raise IndexError("IndexError: invalid index")
        index = index + (slice(None),) * (N - len(index))
        return type(self)(self.model, model_params[index])

    @property
    def shape(self):
        return self.model_params.shape[:-1]

def one_stage(data, gtab, x0, jac, bounds, tol, routine, algorithm,
              gtol, ftol, eps):
    """
    Fit the ivim params using minimize

    Parameters
    ----------
    data : array ([X, Y, Z, ...], g)
        Data or response variables holding the data. Note that the last
        dimension should contain the data. It makes no copies of data.

    jac : bool
        Use the Jacobian? Default: False

    Returns
    -------
    ivim_params: the S0, f, D_star, D value for each voxel.

    """
    flat_data = data.reshape((-1, data.shape[-1]))
    flat_x0 = x0.reshape(-1, x0.shape[-1])
    # Flatten for the iteration over voxels:
    bvals = gtab.bvals
    ivim_params = np.empty((flat_data.shape[0], 4))

    if routine == 'minimize':
        result = _minimize(flat_data, bvals, flat_x0, ivim_params,
                           bounds, tol, jac, algorithm, gtol, ftol, eps)
    elif routine == "leastsq":
        result = _leastsq(flat_data, bvals, flat_x0, ivim_params)
    ivim_params.shape = data.shape[:-1] + (4,)
    return ivim_params

def _minimize(flat_data, bvals, flat_x0, ivim_params,
              bounds, tol, jac, algorithm, gtol, ftol, eps):
    """Use minimize for finding ivim_params"""
    sum_sq = lambda params, bvals, signal: np.sum(
        _ivim_error(params, bvals, signal))

    result = []
    for vox in range(flat_data.shape[0]):
        res = Optimizer(sum_sq,
                        flat_x0[vox],
                        args=(bvals, flat_data[vox]), bounds=bounds,
                        tol=tol, method=algorithm, jac=jac,
                        options={'gtol': gtol, 'ftol': ftol, 'eps': eps})
        ivim_params[vox, :4] = res.xopt
        result += [res]
    return result

def _leastsq(flat_data, bvals, flat_x0, ivim_params):
    """Use minimize for finding ivim_params"""
    result = []
    for vox in range(flat_data.shape[0]):
        res = leastsq(_ivim_error,
                      flat_x0[vox],
                      args=(bvals, flat_data[vox]))
        ivim_params[vox, :4] = res[0]
        result += [res]
    return result

def two_stage(data, gtab, x0,
              split_b, jac, bounds, tol,
              routine, algorithm, gtol, ftol, eps):
    """
    Fit the ivim params using a two stage fit

    Parameters
    ----------
    data : array ([X, Y, Z, ...], g)
        Data or response variables holding the data. Note that the last
        dimension should contain the data. It makes no copies of data.

    jac : bool
        Use the Jacobian? Default: False

    Returns
    -------
    ivim_params: the S0, f, D_star, D value for each voxel.

    """
    flat_data = data.reshape((-1, data.shape[-1]))
    # Flatten for the iteration over voxels:
    bvals_ge_split = gtab.bvals[gtab.bvals > split_b]
    bvecs_ge_split = gtab.bvecs[gtab.bvals > split_b]
    gtab_ge_split = gradient_table(bvals_ge_split, bvecs_ge_split.T)

    tensor_model = TensorModel(gtab_ge_split)
    tenfit = tensor_model.fit(data[..., gtab.bvals > split_b])

    D_guess = mean_diffusivity(tenfit.evals)

    x0[:, 3] = D_guess

    return one_stage(data, gtab, x0, jac, bounds, tol, routine, algorithm,
                     gtol, ftol, eps)
<commit_msg>added f guess<commit_after>#!/usr/bin/python
""" Classes and functions for fitting ivim model """
from __future__ import division, print_function, absolute_import
import warnings
import functools
import numpy as np
from dipy.core.optimize import Optimizer

from scipy.optimize import leastsq

from dipy.core.gradients import gradient_table
from dipy.reconst.base import ReconstModel
from dipy.reconst.dti import _min_positive_signal
from dipy.reconst.dti import TensorModel, mean_diffusivity

def ivim_function(params, bvals):
    """The ivim function
    Parameters
    ----------
    params : array
             The IVIM parameters [S0, f, D_star, D]

    bvals  : array
             bvalues
    """
    S0, f, D_star, D = params
    S = S0 * (f * np.exp(-bvals * D_star) + (1 - f) * np.exp(-bvals * D))
    return S

def _ivim_error(params, bvals, signal):
    """Error function to be used in fitting the model
    """
    return (signal - ivim_function(params, bvals))**2

class IvimModel(ReconstModel):
    """Ivim model
    """

    def __init__(self, gtab, split_b=200.0, min_signal=None):
        """
        An IVIM model

        Parameters
        ----------
        gtab : GradientTable class instance
        split_b : split the data at this b-value for a two stage fit

        References
        ----------
        .. [1] Le Bihan, Denis, et al. "Separation of diffusion
                   and perfusion in intravoxel incoherent motion MR
                   imaging." Radiology 168.2 (1988): 497-505.
        .. [2] Federau, Christian, et al. "Quantitative measurement
                   of brain perfusion with intravoxel incoherent motion
                   MR imaging." Radiology 265.3 (2012): 874-881.
        """
        ReconstModel.__init__(self, gtab)
        self.split_b = split_b
        self.min_signal = min_signal
        if self.min_signal is not None and self.min_signal <= 0:
            e_s = "The `min_signal` key-word argument needs to be strictly"
            e_s += " positive."
            raise ValueError(e_s)

    def fit(self, data, mask=None, x0=None, fit_method="one_stage", routine="minimize",
            jac=False, bounds=((0, 1.), (0, 1.), (0, 1.), (0, 1.)), tol=1e-25, algorithm='L-BFGS-B',
            gtol=1e-25, ftol=1e-25, eps=1e-15):
        """ Fit method of the Ivim model class

        Parameters
        ----------
        data : array
            The measured signal from one voxel.
        mask : array
            A boolean array used to mark the coordinates in the data that
            should be analyzed that has the shape data.shape[:-1]
        x0 : array
            Initial guesses for the parameters S0, f, D_star and D
            Default x0 = [1.0, 0.10, 0.001, 0.0009]
            These guess parameters are taken from the Federau paper
            Dimension can either be 1 or (N, 4) where N is the number of
            voxels in the data.
        fit_method: str
            Use one-stage fitting or two-stage fitting.
        jac : Boolean
            Use the Jacobian
        bounds : tuple
            Bounds for the various parameters
        tol : float
            tolerance for convergence
        method : str
            Fitting algorithm to use from scipy.optimize

        Returns
        -------
        IvimFit object
        """
        if mask is None:
            # Flatten it to 2D either way:
            data_in_mask = np.reshape(data, (-1, data.shape[-1]))
        else:
            # Check for valid shape of the mask
            if mask.shape != data.shape[:-1]:
                raise ValueError("Mask is not the same shape as data.")
            mask = np.array(mask, dtype=bool, copy=False)
            data_in_mask = np.reshape(data[mask], (-1, data.shape[-1]))

        if self.min_signal is None:
            min_signal = _min_positive_signal(data)
        else:
            min_signal = self.min_signal

        # Generate guess_parameters for all voxels
        if x0 is None:
            x0 = np.ones(
                (data.shape[:-1] + (4,))) * [1.0, 0.10, 0.001, 0.0009]
        else:
            if x0.shape != (data.shape[0], 4):
                raise ValueError(
                    "Guess params should be of shape (num_voxels,4)")
        data_in_mask = np.maximum(data_in_mask, min_signal)

        if fit_method == "one_stage":
            params_in_mask = one_stage(data_in_mask, self.gtab, x0,
                                       jac, bounds, tol,
                                       routine,
                                       algorithm, gtol, ftol, eps)
        elif fit_method == "two_stage":
            params_in_mask = two_stage(data_in_mask, self.gtab, x0,
                                       self.split_b, jac, bounds, tol,
                                       routine,
                                       algorithm, gtol, ftol, eps)
        else:
            raise ValueError("""Fit method must be either
                             'one_stage' or 'two_stage'""")

        if mask is None:
            out_shape = data.shape[:-1] + (-1, )
            ivim_params = params_in_mask.reshape(out_shape)
        else:
            ivim_params = np.zeros(data.shape[:-1] + (4,))
            ivim_params[mask, :] = params_in_mask

        return IvimFit(self, ivim_params)

class IvimFit(object):

    def __init__(self, model, model_params):
        """ Initialize a IvimFit class instance.
            Parameters
            ----------
            The model parameters are S0, f, D_star, D
        """
        self.model = model
        self.model_params = model_params

    def __getitem__(self, index):
        model_params = self.model_params
        N = model_params.ndim
        if type(index) is not tuple:
            index = (index,)
        elif len(index) >= model_params.ndim:
            raise IndexError("IndexError: invalid index")
        index = index + (slice(None),) * (N - len(index))
        return type(self)(self.model, model_params[index])

    @property
    def shape(self):
        return self.model_params.shape[:-1]

def one_stage(data, gtab, x0, jac, bounds, tol, routine, algorithm,
              gtol, ftol, eps):
    """
    Fit the ivim params using minimize

    Parameters
    ----------
    data : array ([X, Y, Z, ...], g)
        Data or response variables holding the data. Note that the last
        dimension should contain the data. It makes no copies of data.

    jac : bool
        Use the Jacobian? Default: False

    Returns
    -------
    ivim_params: the S0, f, D_star, D value for each voxel.

    """
    flat_data = data.reshape((-1, data.shape[-1]))
    flat_x0 = x0.reshape(-1, x0.shape[-1])
    # Flatten for the iteration over voxels:
    bvals = gtab.bvals
    ivim_params = np.empty((flat_data.shape[0], 4))

    if routine == 'minimize':
        result = _minimize(flat_data, bvals, flat_x0, ivim_params,
                           bounds, tol, jac, algorithm, gtol, ftol, eps)
    elif routine == "leastsq":
        result = _leastsq(flat_data, bvals, flat_x0, ivim_params)
    ivim_params.shape = data.shape[:-1] + (4,)
    return ivim_params

def _minimize(flat_data, bvals, flat_x0, ivim_params,
              bounds, tol, jac, algorithm, gtol, ftol, eps):
    """Use minimize for finding ivim_params"""
    sum_sq = lambda params, bvals, signal: np.sum(
        _ivim_error(params, bvals, signal))

    result = []
    for vox in range(flat_data.shape[0]):
        res = Optimizer(sum_sq,
                        flat_x0[vox],
                        args=(bvals, flat_data[vox]), bounds=bounds,
                        tol=tol, method=algorithm, jac=jac,
                        options={'gtol': gtol, 'ftol': ftol, 'eps': eps})
        ivim_params[vox, :4] = res.xopt
        result += [res]
    return result

def _leastsq(flat_data, bvals, flat_x0, ivim_params):
    """Use minimize for finding ivim_params"""
    result = []
    for vox in range(flat_data.shape[0]):
        res = leastsq(_ivim_error,
                      flat_x0[vox],
                      args=(bvals, flat_data[vox]))
        ivim_params[vox, :4] = res[0]
        result += [res]
    return result

def two_stage(data, gtab, x0,
              split_b, jac, bounds, tol,
              routine, algorithm, gtol, ftol, eps):
    """
    Fit the ivim params using a two stage fit

    Parameters
    ----------
    data : array ([X, Y, Z, ...], g)
        Data or response variables holding the data. Note that the last
        dimension should contain the data. It makes no copies of data.

    jac : bool
        Use the Jacobian? Default: False

    Returns
    -------
    ivim_params: the S0, f, D_star, D value for each voxel.

    """
    flat_data = data.reshape((-1, data.shape[-1]))
    # Flatten for the iteration over voxels:
    bvals_ge_split = gtab.bvals[gtab.bvals > split_b]
    bvecs_ge_split = gtab.bvecs[gtab.bvals > split_b]
    gtab_ge_split = gradient_table(bvals_ge_split, bvecs_ge_split.T)

    tensor_model = TensorModel(gtab_ge_split)
    tenfit = tensor_model.fit(data[..., gtab.bvals > split_b])

    D_guess = mean_diffusivity(tenfit.evals)

    x0[..., 3] = D_guess

    # Calculate the intercept for straight line considering bvals > split_b
    # This will be a straight line with slope D_guess(m). If y = mx + C is the line
    # we can get the intercept by putting x = split_b, y = data(split_b)
    # Thus, C = data(b = split_b) - D_guess*split_b

    # The guess for f is given by 1 - S0/C, where S0 is the data(b = 0) value

    C = data[..., gtab.bvals == split_b][0] - D_guess * split_b
    f_guess = data[..., 0] / C
    x0[..., 1] = f_guess

    return one_stage(data, gtab, x0, jac, bounds, tol, routine, algorithm,
                     gtol, ftol, eps)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import os
import re
import yaml
import click
import importlib
from functools import partial

# Module API

@click.command()
@click.argument('path', default='.')
def cli(path):
    specs = parse_specs(path)
    success = test_specs(specs)
    if not success:
        exit(1)

# Helpers

def parse_specs(path):
    # Maps
    specmap = {}
    hookmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if not filename.endswith('.yml') and filename != 'packspec.py':
                continue
            filepath = os.path.join(root, filename)
            filecont = io.open(filepath, encoding='utf-8').read()
            if filename == 'packspec.py':
                scope = {}
                exec(filecont, scope)
                for name, attr in scope.items():
                    if name.startswith('_'):
                        continue
                    hookmap[name] = attr
                continue
            spec = parse_spec(filecont)
            if not spec:
                continue
            if spec['scope']['PACKAGE'] not in specmap:
                specmap[spec['scope']['PACKAGE']] = spec
            else:
                specmap[spec['scope']['PACKAGE']]['features'].extend(spec['features'])
    # Specs
    specs = [specmap[package] for package in sorted(specmap)]
    for spec in specs:
        for name, hook in hookmap.items():
            spec['scope'][name] = partial(hook, spec['scope'])
    return specs

def parse_spec(spec):
    # Package
    contents = yaml.load(spec)
    try:
        feature = parse_feature(contents[0])
        package = feature['result']
        assert feature['property'] == 'PACKAGE'
        assert not feature['skip']
    except Exception:
        return None
    # Features
    features = []
    for feature in contents:
        feature = parse_feature(feature)
        features.append(feature)
    # Scope
    scope = {'PACKAGE': package}
    module = importlib.import_module(package)
    for name in dir(module):
        if name.startswith('_'):
            continue
        scope[name] = getattr(module, name)
    return {
        'features': features,
        'scope': scope,
    }

def parse_feature(feature):
    left, right = list(feature.items())[0]
    # Left side
    match = re.match(r'^(?:([^=]*)=)?([^:]*)(?::(.*))*$', left)
    assign, property, skip = match.groups()
    if skip:
        filters = skip.split(':')
        skip = '!py' in filters or not ('!' in skip or 'py' in filters)
    # Right side
    result = right
    arguments = None
    if isinstance(right, list):
        result = right[-1]
        arguments = right[:-1]
    # Text repr
    text = property
    if assign:
        text = '%s=%s' % (assign, text)
    if arguments:
        text = '%s(%s)' % (text, ', '.join(map(repr, arguments)))
    if not assign:
        text = '%s == %s' % (text, repr(result))
    return {
        'assign': assign,
        'property': property,
        'arguments': arguments,
        'result': result,
        'text': text,
        'skip': skip,
    }

def test_specs(specs):
    success = True
    for spec in specs:
        spec_success = test_spec(spec)
        success = success and spec_success
    return success

def test_spec(spec):
    passed = 0
    amount = len(spec['features'])
    for feature in spec['features']:
        passed += test_feature(feature, spec['scope'])
    print('%s: %s/%s' % (spec['scope']['PACKAGE'], passed, amount))
    success = (passed == amount)
    return success

def test_feature(feature, scope):
    # Skip
    if feature['skip']:
        print('(#) %s' % feature['text'])
        return True
    # Execute
    try:
        owner = scope
        parts = feature['property'].split('.')
        for name in parts[:-1]:
            owner = get_property(owner, name)
        property = get_property(owner, parts[-1])
        # Call property
        if feature['arguments'] is not None:
            arguments = []
            for argument in feature['arguments']:
                arguments.append(argument)
            result = property(*arguments)
        # Get property
        elif property is not None:
            result = property
        # Set property
        else:
            result = feature['result']
            set_property(owner, parts[-1], result)
    except Exception:
        result = 'ERROR'
    # Assign
    if feature['assign'] is not None:
        scope[feature['assign']] = result
    # Verify
    success = result == feature['result'] or (result != 'ERROR' and feature['result'] == 'ANY')
    if success:
        print('(+) %s' % feature['text'])
    else:
        print('(-) %s # %s' % (feature['text'], repr(result)))
    return success

def get_property(owner, name):
    if isinstance(owner, dict):
        return owner.get(name)
    return getattr(owner, name, None)

def set_property(owner, name, value):
    if isinstance(owner, dict):
        owner[name] = value
        return
    return setter(owner, name, value)

# Main program

if __name__ == '__main__':
    cli()
<commit_msg>implemented variable interpolation on right side<commit_after># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import os
import re
import yaml
import click
import importlib
from functools import partial

# Module API

@click.command()
@click.argument('path', default='.')
def cli(path):
    specs = parse_specs(path)
    success = test_specs(specs)
    if not success:
        exit(1)

# Helpers

def parse_specs(path):
    # Maps
    specmap = {}
    hookmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if not filename.endswith('.yml') and filename != 'packspec.py':
                continue
            filepath = os.path.join(root, filename)
            filecont = io.open(filepath, encoding='utf-8').read()
            if filename == 'packspec.py':
                scope = {}
                exec(filecont, scope)
                for name, attr in scope.items():
                    if name.startswith('_'):
                        continue
                    hookmap[name] = attr
                continue
            spec = parse_spec(filecont)
            if not spec:
                continue
            if spec['scope']['PACKAGE'] not in specmap:
                specmap[spec['scope']['PACKAGE']] = spec
            else:
                specmap[spec['scope']['PACKAGE']]['features'].extend(spec['features'])
    # Specs
    specs = [specmap[package] for package in sorted(specmap)]
    for spec in specs:
        for name, hook in hookmap.items():
            spec['scope'][name] = partial(hook, spec['scope'])
    return specs

def parse_spec(spec):
    # Package
    contents = yaml.load(spec)
    try:
        feature = parse_feature(contents[0])
        package = feature['result']
        assert feature['property'] == 'PACKAGE'
        assert not feature['skip']
    except Exception:
        return None
    # Features
    features = []
    for feature in contents:
        feature = parse_feature(feature)
        features.append(feature)
    # Scope
    scope = {'PACKAGE': package}
    module = importlib.import_module(package)
    for name in dir(module):
        if name.startswith('_'):
            continue
        scope[name] = getattr(module, name)
    return {
        'features': features,
        'scope': scope,
    }

def parse_feature(feature):
    left, right = list(feature.items())[0]
    # Left side
    match = re.match(r'^(?:([^=]*)=)?([^:]*)(?::(.*))*$', left)
    assign, property, skip = match.groups()
    if skip:
        filters = skip.split(':')
        skip = '!py' in filters or not ('!' in skip or 'py' in filters)
    # Right side
    result = right
    arguments = None
    if isinstance(right, list):
        result = right[-1]
        arguments = right[:-1]
    # Text repr
    text = property
    if assign:
        text = '%s=%s' % (assign, text)
    if arguments is not None:
        items = []
        for argument in arguments:
            item = parse_interpolation(argument)
            if not item:
                item = repr(argument)
            items.append(item)
        text = '%s(%s)' % (text, ', '.join(items))
    if not assign:
        text = '%s == %s' % (text, repr(result))
    return {
        'assign': assign,
        'property': property,
        'arguments': arguments,
        'result': result,
        'text': text,
        'skip': skip,
    }

def test_specs(specs):
    success = True
    for spec in specs:
        spec_success = test_spec(spec)
        success = success and spec_success
    return success

def test_spec(spec):
    passed = 0
    amount = len(spec['features'])
    for feature in spec['features']:
        passed += test_feature(feature, spec['scope'])
    print('%s: %s/%s' % (spec['scope']['PACKAGE'], passed, amount))
    success = (passed == amount)
    return success

def test_feature(feature, scope):
    # Skip
    if feature['skip']:
        print('(#) %s' % feature['text'])
        return True
    # Execute
    try:
        owner = scope
        names = feature['property'].split('.')
        for name in names[:-1]:
            owner = get_property(owner, name)
        property = get_property(owner, names[-1])
        # Call property
        if feature['arguments'] is not None:
            arguments = []
            for argument in feature['arguments']:
                # Property interpolation
                name = parse_interpolation(argument)
                if name:
                    argument = get_property(scope, name)
                arguments.append(argument)
            result = property(*arguments)
        # Get property
        elif property is not None:
            result = property
        # Set property
        else:
            result = feature['result']
            set_property(owner, names[-1], result)
    except Exception:
        result = 'ERROR'
    # Assign
    if feature['assign'] is not None:
        scope[feature['assign']] = result
    # Verify
    success = result == feature['result'] or (result != 'ERROR' and feature['result'] == 'ANY')
    if success:
        print('(+) %s' % feature['text'])
    else:
        print('(-) %s # %s' % (feature['text'], repr(result)))
    return success

def parse_interpolation(argument):
    if isinstance(argument, dict) and len(argument) == 1:
        left, right = list(argument.items())[0]
        if right is None:
            return left
    return None

def get_property(owner, name):
    if isinstance(owner, dict):
        return owner.get(name)
    return getattr(owner, name, None)

def set_property(owner, name, value):
    if isinstance(owner, dict):
        owner[name] = value
        return
    return setter(owner, name, value)

# Main program

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


<commit_before>import numpy as np

class ion_datacube():
    # a class for holding a datacube from an MSI dataset and producing ion images
    def __init__(self, step_size=[]):  # define mandatory parameters (todo- validator to check all present)
        self.xic = []  # 'xtracted ion chromatogram (2d array of intensities(vector of time by mzs))
        self.bounding_box = []  # tuple (x,y,w,h) in co-ordinate space
        self.coords = []  # co-ordiantes of every pixel
        self.mzs = []  # centroids of each xic vector
        self.tol = []  # tolerance window around centroids in m/z
        if step_size == []:
            self.step_size = []
        else:
            self.step_size = step_size
        self.n_im = 0

    def add_xic(self, xic, mz, tol):
        # add an eXtracted Ion Chromatogram to the datacube
        if len(self.coords) != len(xic):
            raise ValueError(
                'size of co-ordinates to not match xic (coords:{} xic:{})'.format(len(self.coords), len(xic)))
        self.xic.append(xic)  # = np.concatenate((self.xic,xic),axis=1)
        self.mzs = np.concatenate((self.mzs, mz))
        self.tol = np.concatenate((self.tol, tol))
        self.n_im += 1

    def remove_xic(self, mz):
        # remove an xic and related info
        raise NotImplementedError
        index_to_remove = self.mz.index(mz)
        self.n_im -= 1

    def add_coords(self, coords):
        # record spatial co-ordinates for each spectrum
        # coords is a ix3 array with x,y,z coords
        self.coords = coords
        # if len(self.xic)==0:
        #    self.xic = np.zeros((len(coords),0))
        self.calculate_bounding_box()
        self.coord_to_index()

    def calculate_bounding_box(self):
        self.bounding_box = []
        for ii in range(0, 3):
            self.bounding_box.append(np.amin(self.coords[:, ii]))
            self.bounding_box.append(np.amax(self.coords[:, ii]))

    def coord_to_index(self, transform_type='reg_grid', params=()):
        # this function maps coords onto pixel indicies (assuming a grid defined by bounding box and transform type)
        # -implement methods such as interp onto grid spaced over coords
        # -currently treats coords as grid positions,
        if transform_type == 'reg_grid':
            # data was collected on a grid
            # - coordinates can transformed directly into pixel indiceis
            # - subtract smallest value and divide by x & y step size
            pixel_indices = np.zeros(len(self.coords))
            _coord = np.asarray(self.coords)
            _coord = np.around(_coord, 5)  # correct for numerical precision
            _coord = _coord - np.amin(_coord, axis=0)
            if self.step_size == []:  # no additional info, guess step size in xyz
                self.step_size = np.zeros((3, 1))
                for ii in range(0, 3):
                    self.step_size[ii] = np.mean(np.diff(np.unique(_coord[:, ii])))

            # coordinate to pixels
            _coord /= np.reshape(self.step_size, (3,))
            _coord_max = np.amax(_coord, axis=0)
            self.nColumns = _coord_max[1] + 1
            self.nRows = _coord_max[0] + 1
            pixel_indices = _coord[:, 0] * self.nColumns + _coord[:, 1]
            pixel_indices = pixel_indices.astype(np.int32)
        else:
            print 'transform type not recognised'
            raise ValueError
        self.pixel_indices = pixel_indices

    def xic_to_image(self, xic_index):
        xic = self.xic[xic_index].copy()
        # turn xic into an image
        img = -1 + np.zeros(self.nRows * self.nColumns)
        img[self.pixel_indices] = xic
        img = np.reshape(img, (self.nRows, self.nColumns))
        return img

    def xic_to_image_fixed_size(self, xic_index, out_im_size):
        def calc_pad(px1, px2):
            delta = px1-px2
            if delta < 0:
                return 0, 0
            return delta/2, delta/2+delta%2

        def calc_crop(px1, px2):
            delta = px1 - px2
            if delta > 0:
                return 0, -1
            delta = abs(delta)
            return delta/2, -1*delta/2+delta%2

        pad_total_pixels = map(calc_pad, out_im_size, (self.nRows, self.nColumns))
        crop_total_pixels = map(calc_crop, out_im_size, (self.nRows, self.nColumns))

        img = self.xic_to_image(xic_index)
        img=np.asarray(img)
        img = img[crop_total_pixels[0][0]:crop_total_pixels[0][1], crop_total_pixels[1][0]:crop_total_pixels[1][1]]
        img = np.pad(img, ((pad_total_pixels[0][0], pad_total_pixels[0][1]), (pad_total_pixels[1][0], pad_total_pixels[1][1])), 'constant', constant_values=-1)
        return img

    def image_to_xic(self, im):
        """
        takes an image (presumably generated by xic to image) and returns a vector suitable for insertion into the datacube
        :param im:
            numpy 2d array
        :return:
            numpy vector
        """
        im = np.reshape(im, self.nRows * self.nColumns)
        xic = im[self.pixel_indices]
        return xic

    def apply_image_processing(self, smooth_method, **smooth_params):
        """
        Function to apply pre-defined image processing methods to ion_datacube
        #todo: expose parameters in config
        :param ion_datacube:
            object from pyImagingMSpec.ion_datacube already containing images
        :return:
            ion_datacube is updated in place.
            None returned
        """
        from pyImagingMSpec import smoothing
        for ii in range(self.n_im):
            im = self.xic_to_image(ii)
            # todo: for method in smoothing_methods:
            methodToCall = getattr(smoothing, smooth_method)
            im_s = methodToCall(im, **smooth_params)
            self.xic[ii] = self.image_to_xic(im_s)
<commit_msg>updated comments<commit_after>import numpy as np

class ion_datacube():
    # a class for holding a datacube from an MSI dataset and producing ion images
    def __init__(self, step_size=[]):  # define mandatory parameters (todo- validator to check all present)
        self.xic = []  # 'xtracted ion chromatogram (2d array of intensities(vector of time by mzs))
        self.bounding_box = []  # tuple (x,y,w,h) in co-ordinate space
        self.coords = []  # co-ordiantes of every pixel
        self.mzs = []  # centroids of each xic vector
        self.tol = []  # tolerance window around centroids in m/z
        if step_size == []:
            self.step_size = []
        else:
            self.step_size = step_size
        self.n_im = 0

    def add_xic(self, xic, mz, tol):
        # add an eXtracted Ion Chromatogram to the datacube
        if len(self.coords) != len(xic):
            raise ValueError(
                'size of co-ordinates to not match xic (coords:{} xic:{})'.format(len(self.coords), len(xic)))
        self.xic.append(xic)  # = np.concatenate((self.xic,xic),axis=1)
        self.mzs = np.concatenate((self.mzs, mz))
        self.tol = np.concatenate((self.tol, tol))
        self.n_im += 1

    def remove_xic(self, mz):
        # remove an xic and related info
        raise NotImplementedError
        index_to_remove = self.mz.index(mz)
        self.n_im -= 1

    def add_coords(self, coords):
        # record spatial co-ordinates for each spectrum
        # coords is a ix3 array with x,y,z coords
        self.coords = coords
        # if len(self.xic)==0:
        #    self.xic = np.zeros((len(coords),0))
        self.calculate_bounding_box()
        self.coord_to_index()

    def calculate_bounding_box(self):
        self.bounding_box = []
        for ii in range(0, 3):
            self.bounding_box.append(np.amin(self.coords[:, ii]))
            self.bounding_box.append(np.amax(self.coords[:, ii]))

    def coord_to_index(self, transform_type='reg_grid', params=()):
        # this function maps coords onto pixel indicies (assuming a grid defined by bounding box and transform type)
        # -implement methods such as interp onto grid spaced over coords
        # -currently treats coords as grid positions,
        if transform_type == 'reg_grid':
            # data was collected on a grid
            # - coordinates can transformed directly into pixel indiceis
            # - subtract smallest value and divide by x & y step size
            pixel_indices = np.zeros(len(self.coords))
            _coord = np.asarray(self.coords)
            _coord = np.around(_coord, 5)  # correct for numerical precision
            _coord = _coord - np.amin(_coord, axis=0)
            if self.step_size == []:  # no additional info, guess step size in xyz
                self.step_size = np.zeros((3, 1))
                for ii in range(0, 3):
                    self.step_size[ii] = np.mean(np.diff(np.unique(_coord[:, ii])))

            # coordinate to pixels
            _coord /= np.reshape(self.step_size, (3,))
            _coord_max = np.amax(_coord, axis=0)
            self.nColumns = _coord_max[1] + 1
            self.nRows = _coord_max[0] + 1
            pixel_indices = _coord[:, 0] * self.nColumns + _coord[:, 1]
            pixel_indices = pixel_indices.astype(np.int32)
        else:
            print 'transform type not recognised'
            raise ValueError
        self.pixel_indices = pixel_indices

    def xic_to_image(self, xic_index):
        xic = self.xic[xic_index].copy()
        # turn xic into an image
        img = -1 + np.zeros(self.nRows * self.nColumns)
        img[self.pixel_indices] = xic
        img = np.reshape(img, (self.nRows, self.nColumns))
        return img

    def xic_to_image_fixed_size(self, xic_index, out_im_size):
        """
        returns an xic formatted as an image. This function returns an image of a pre-specfied size, cropping or padding equally around the image as required.
        :param xic_index: index of the xic to turn into an image
        :param out_im_size: desired image size
        :return:
        """
        def calc_pad(px1, px2):
            delta = px1-px2
            if delta < 0:
                return 0, 0
            return delta/2, delta/2+delta%2

        def calc_crop(px1, px2):
            delta = px1 - px2
            if delta > 0:
                return 0, -1
            delta = abs(delta)
            return delta/2, -1*delta/2+delta%2

        pad_total_pixels = map(calc_pad, out_im_size, (self.nRows, self.nColumns))
        crop_total_pixels = map(calc_crop, out_im_size, (self.nRows, self.nColumns))

        img = self.xic_to_image(xic_index)
        img=np.asarray(img)
        img = img[crop_total_pixels[0][0]:crop_total_pixels[0][1], crop_total_pixels[1][0]:crop_total_pixels[1][1]]
        img = np.pad(img, ((pad_total_pixels[0][0], pad_total_pixels[0][1]), (pad_total_pixels[1][0], pad_total_pixels[1][1])), 'constant', constant_values=-1)
        return img

    def image_to_xic(self, im):
        """
        takes an image (presumably generated by xic to image) and returns a vector suitable for insertion into the datacube
        :param im:
            numpy 2d array
        :return:
            numpy vector
        """
        im = np.reshape(im, self.nRows * self.nColumns)
        xic = im[self.pixel_indices]
        return xic

    def apply_image_processing(self, smooth_method, **smooth_params):
        """
        Function to apply pre-defined image processing methods to ion_datacube
        #todo: expose parameters in config
        :param ion_datacube:
            object from pyImagingMSpec.ion_datacube already containing images
        :return:
            ion_datacube is updated in place.
            None returned
        """
        from pyImagingMSpec import smoothing
        for ii in range(self.n_im):
            im = self.xic_to_image(ii)
            # todo: for method in smoothing_methods:
            methodToCall = getattr(smoothing, smooth_method)
            im_s = methodToCall(im, **smooth_params)
            self.xic[ii] = self.image_to_xic(im_s)
<|endoftext|>


<commit_before>__version__ = '0.1.4'

default_app_config = 'connected_accounts.apps.ConnectedAccountsConfig'
<commit_msg>Bump version.<commit_after>__version__ = '0.1.5'

default_app_config = 'connected_accounts.apps.ConnectedAccountsConfig'
<|endoftext|>


<commit_before>import logging
import sys
import subprocess
import csv
import re
import datetime
from . import *
from util import zfs_list, is_syncing, zfs_destroy, zfs_snapshot

PREFIX="zfs-auto-snap"
USERPROP_NAME='com.sun:auto-snapshot'
SEP=":"
KEEP={'hourly': 24, 'daily': 30, '__default__': 10}

class AutoSnapshotter():
    """Automatically snapshot ZFS filesystems

    This class will manage automatic snapshots for ZFS
    filesystems. Snapshots will be named according to the
    pattern #{prefix}_#{label}-#{date}.

    Older snapshots for the given label value will be
    automatically purged, with the most recent items
    retained per the keep attribute. If keep is the special
    value 'all', all older snapshots are retained.

    Attributes:
        label           The label for this set of snapshots
        keep            Number of older snapshots to keep, or 'all'
        avoidsync       Avoid fs on zpools in scrub/resilver state
        prefix          First part of snapshot name
        userprop_name   name of the ZFS user property to check
    """
    def __init__(
        self,
        label,
        keep='all',
        avoidsync=False,
        prefix=PREFIX,
        userprop_name=USERPROP_NAME
    ):
        """Create new AutoSnapshotter instance
        """

        self.keep          = keep
        self.label         = label
        self.avoidsync     = avoidsync
        self.prefix        = prefix
        self.userprop_name = userprop_name

    def take_snapshot(self, fsnames, snap_children=False):
        """Take a snapshot of all eligible filesystems given in fsnames

        If fsnames is the special value '//', use the #{self.userprop_name} or
        the #{self.userprop_name}:#{self.label} property to determine which
        filesystems to snapshot
        """

        today = datetime.datetime.now()
        snapdate=today.strftime('%F-%H%M')
        logging.debug("Snapdate is: %s"% snapdate)
        snapname="%s_%s-%s" % (self.prefix, self.label, snapdate)

        # the '//' filesystem is special. We use it as a keyword to determine
        # whether to poll the ZFS user properties.
        # Determine what these are, call ourselves again, then return.
        if isinstance(fsnames, basestring) and fsnames == '//':
            single_list,recursive_list = get_userprop_datasets(
                label=self.label, userprop_name=self.userprop_name)

            logging.info("Taking non-recursive snapshots of: %s" %\
                           ', '.join(single_list))
            single_state = self.take_snapshot(single_list, snap_children=False)

            logging.info("Taking recursive snapshots of: %s" %\
                           ', '.join(recursive_list))
            recursive_state = self.take_snapshot(
                recursive_list, snap_children=True)

            return

        if self.avoidsync == True:
            fsnames=filter_syncing_pools(fsnames)

        keep=self.keep
        # Since we are about to take a new snapshot, get rid of 1 extra
        if keep != 'all':
            keep = keep - 1

        for fs in fsnames:
            # Ok, now say cheese!
            logging.info("Taking %s snapshot %s@%s" % (
                "recursive" if snap_children else "non-recursive",
                fs,
                snapname))
            zfs_snapshot(fs,snapname,snap_children)

            # If we're taking recursive snapshots,
            # walk through the children, destroying old ones if required.
            destroy_older_snapshots(fs, keep, self.label,
                                    self.prefix, snap_children)
        pass

def destroy_older_snapshots(filesys, keep, label, prefix=PREFIX,
                            recursive=False):
    """Destroy old snapshots, keeping 'keep' newest around.

    Given a filesystem name, the number of snapshots we want to keep, along
    with the label for this set of snapshots, we destroy all older
    snapshots of this filesystem whose names being with the text
    prefix_label-

    If keep is set to the special value 'all', no older snapshots are
    removed.

    Note that unlike the original ksh function, we actually keep around the
    requested number of snapshots, rather than "keep - 1".

    Returns the number of snapshots removed.
    """

    if keep == 'all':
        return 0

    snappre="%s@%s_%s-" % (filesys, prefix, label)
    r = zfs_list(types=['snapshot'], sort='creation', properties=['name'],
                 ds=filesys, recursive=True)

    # Remove all snapshots for child filesystems and those that aren't for
    # our given label
    rs = [x for x in r if x[:len(snappre)] == snappre]

    to_remove=rs[keep:]
    removed=0
    logging.debug(
        "Should remove %d of %d snapshots for filesys %s (keep=%d)" % (
        len(to_remove), len(rs), filesys, keep))
    for snapshot in to_remove:
        try:
            zfs_destroy(snapshot, recursive=recursive)
        # Not catching ZfsArgumentErrors because those are programming problems
        except (ZfsOSError) as e:
            logger.warning('Unable to destroy %s' % snapshot)
        else:
            removed+=1

    return removed

def filter_syncing_pools(fsnames):
    """filter out filesys on pools that are scrubbing/resilvering

    Given a list of fsnames, filter out the filesystems that are on pools
    in the middle of a sync operating (scrub or resilver).

    Sync operations would be interrupted/restarted by a snapshot.

    This function attempts to cache calls to zfs.utils.is_syncing. There is a
    risk that a scrub/resilver will be started just after this check completes,
    and also the risk that a running scrub will complete just after this check.
    """

    poolcache={}
    nosyncfilesys=[]

    for fs in fsnames:
        pool=fs.split('/')[0]
        if pool == '':
            raise ZfsBadFsName(fs)

        if pool in poolcache.keys():
            if poolcache[pool]==True:
                sync=True
            else:
                sync=False
        else:
            sync = is_syncing(pool)
            poolcache[pool]=sync

        if sync==True:
            logging.info("The pool containing %s is being scrubbed/resilvered." % fs)
            logging.info("Not taking snapshots for %s." % fs)
        else:
            nosyncfilesys.append(fs)
    return nosyncfilesys

def can_recursive_snapshot(ds, excludes):
    """ Given a list of datasets to exclude, check if the dataset named ds
    has a child filesystem that is in the exclude list. If so, ds cannot be
    recursively snapshotted. """
    sds=ds.split('/')
    for exc in excludes:
        sexc=exc.split('/')
        if len(sds) > len(sexc):
            next
        if sds == sexc[:len(sds)]:
            return False
    return True

def narrow_recursive_filesystems(recursive_list):
    final_list=[]
    for ds in recursive_list:
        sds=ds.split('/')
        found=False
        logging.debug("checking if %s has a parent in %s" % (ds,recursive_list))
        for tmp in recursive_list:
            logging.debug("comparing %s to %s" % (tmp,ds))
            if tmp!=ds:
                stmp=tmp.split('/')
                if stmp == sds[:len(stmp)]:
                    found=True
                    next
        if found==False:
            final_list.append(ds)
    return final_list

def get_userprop_datasets(label="daily", userprop_name=USERPROP_NAME):
    """ This builds two lists of datasets - RECURSIVE_LIST and SINGLE_LIST
    based on the value of ZFS user properties com.sun:auto-snapshot and
    com.sun:auto-snapshot:${label}
    RECURSIVE_LIST is a list of datasets that can be snapshotted with -r
    SINGLE_LIST is a list of datasets to snapshot individually.
    """

    props=['name',
           userprop_name,
           SEP.join([userprop_name,label]) ]

    r=zfs_list(sort='name', properties=props)
    save=[]
    exclude=[]
    for row in r:
        save.append(row)
        if row[2] == 'false' or (row[1] == 'false' and row[2] == '-') or \
           (row[1] == '-' and row[2] == '-'):
            exclude.append(row[0])

    recursive_list=[]
    single_list=[]

    # find recursive snapshot sets
    for row in save:
        ds=row[0]
        logging.debug("checking %s" % ds)
        if can_recursive_snapshot(ds,exclude):
            logging.debug("OK to recursive snapshot %s" % ds)
            recursive_list.append(ds)
        elif ds not in exclude:
            logging.debug("OK to snapshot sole dataset %s" % ds)
            single_list.append(ds)
        else:
            logging.debug("NOT OK to snapshot %s" % ds)

    logging.debug("Pre-Narrowed recursive list is %s" % recursive_list)
    final_recursive_list=narrow_recursive_filesystems(recursive_list)
    logging.debug("Final recursive list is %s" % final_recursive_list)
    return (single_list,final_recursive_list)

<commit_msg>Add more debug output<commit_after>import logging
import sys
import subprocess
import csv
import re
import datetime
from . import *
from util import zfs_list, is_syncing, zfs_destroy, zfs_snapshot

PREFIX="zfs-auto-snap"
USERPROP_NAME='com.sun:auto-snapshot'
SEP=":"
KEEP={'hourly': 24, 'daily': 30, '__default__': 10}

class AutoSnapshotter():
    """Automatically snapshot ZFS filesystems

    This class will manage automatic snapshots for ZFS
    filesystems. Snapshots will be named according to the
    pattern #{prefix}_#{label}-#{date}.

    Older snapshots for the given label value will be
    automatically purged, with the most recent items
    retained per the keep attribute. If keep is the special
    value 'all', all older snapshots are retained.

    Attributes:
        label           The label for this set of snapshots
        keep            Number of older snapshots to keep, or 'all'
        avoidsync       Avoid fs on zpools in scrub/resilver state
        prefix          First part of snapshot name
        userprop_name   name of the ZFS user property to check
    """
    def __init__(
        self,
        label,
        keep='all',
        avoidsync=False,
        prefix=PREFIX,
        userprop_name=USERPROP_NAME
    ):
        """Create new AutoSnapshotter instance
        """

        self.keep          = keep
        self.label         = label
        self.avoidsync     = avoidsync
        self.prefix        = prefix
        self.userprop_name = userprop_name

    def take_snapshot(self, fsnames, snap_children=False):
        """Take a snapshot of all eligible filesystems given in fsnames

        If fsnames is the special value '//', use the #{self.userprop_name} or
        the #{self.userprop_name}:#{self.label} property to determine which
        filesystems to snapshot
        """

        today = datetime.datetime.now()
        snapdate=today.strftime('%F-%H%M')
        logging.debug("Snapdate is: %s"% snapdate)
        snapname="%s_%s-%s" % (self.prefix, self.label, snapdate)

        # the '//' filesystem is special. We use it as a keyword to determine
        # whether to poll the ZFS user properties.
        # Determine what these are, call ourselves again, then return.
        if isinstance(fsnames, basestring) and fsnames == '//':
            single_list,recursive_list = get_userprop_datasets(
                label=self.label, userprop_name=self.userprop_name)

            logging.info("Taking non-recursive snapshots of: %s" %\
                           ', '.join(single_list))
            single_state = self.take_snapshot(single_list, snap_children=False)

            logging.info("Taking recursive snapshots of: %s" %\
                           ', '.join(recursive_list))
            recursive_state = self.take_snapshot(
                recursive_list, snap_children=True)

            return

        if self.avoidsync == True:
            fsnames=filter_syncing_pools(fsnames)

        keep=self.keep
        # Since we are about to take a new snapshot, get rid of 1 extra
        if keep != 'all':
            keep = keep - 1

        for fs in fsnames:
            # Ok, now say cheese!
            logging.info("Taking %s snapshot %s@%s" % (
                "recursive" if snap_children else "non-recursive",
                fs,
                snapname))
            zfs_snapshot(fs,snapname,snap_children)

            # If we're taking recursive snapshots,
            # walk through the children, destroying old ones if required.
            destroy_older_snapshots(fs, keep, self.label,
                                    self.prefix, snap_children)
        pass

def destroy_older_snapshots(filesys, keep, label, prefix=PREFIX,
                            recursive=False):
    """Destroy old snapshots, keeping 'keep' newest around.

    Given a filesystem name, the number of snapshots we want to keep, along
    with the label for this set of snapshots, we destroy all older
    snapshots of this filesystem whose names being with the text
    prefix_label-

    If keep is set to the special value 'all', no older snapshots are
    removed.

    Note that unlike the original ksh function, we actually keep around the
    requested number of snapshots, rather than "keep - 1".

    Returns the number of snapshots removed.
    """

    if keep == 'all':
        return 0

    snappre="%s@%s_%s-" % (filesys, prefix, label)
    r = zfs_list(types=['snapshot'], sort='creation', properties=['name'],
                 ds=filesys, recursive=True)

    # Remove all snapshots for child filesystems and those that aren't for
    # our given label
    rs = [x for x in r if x[:len(snappre)] == snappre]

    logging.debug("Got back %d total snapshots pre subset, %d post" % (
        len(r), len(rs))
    to_remove=rs[keep:]
    removed=0
    logging.debug(
        "Should remove %d of %d snapshots for filesys %s (keep=%d)" % (
        len(to_remove), len(rs), filesys, keep))
    for snapshot in to_remove:
        try:
            zfs_destroy(snapshot, recursive=recursive)
        # Not catching ZfsArgumentErrors because those are programming problems
        except (ZfsOSError) as e:
            logger.warning('Unable to destroy %s' % snapshot)
        else:
            removed+=1

    return removed

def filter_syncing_pools(fsnames):
    """filter out filesys on pools that are scrubbing/resilvering

    Given a list of fsnames, filter out the filesystems that are on pools
    in the middle of a sync operating (scrub or resilver).

    Sync operations would be interrupted/restarted by a snapshot.

    This function attempts to cache calls to zfs.utils.is_syncing. There is a
    risk that a scrub/resilver will be started just after this check completes,
    and also the risk that a running scrub will complete just after this check.
    """

    poolcache={}
    nosyncfilesys=[]

    for fs in fsnames:
        pool=fs.split('/')[0]
        if pool == '':
            raise ZfsBadFsName(fs)

        if pool in poolcache.keys():
            if poolcache[pool]==True:
                sync=True
            else:
                sync=False
        else:
            sync = is_syncing(pool)
            poolcache[pool]=sync

        if sync==True:
            logging.info("The pool containing %s is being scrubbed/resilvered." % fs)
            logging.info("Not taking snapshots for %s." % fs)
        else:
            nosyncfilesys.append(fs)
    return nosyncfilesys

def can_recursive_snapshot(ds, excludes):
    """ Given a list of datasets to exclude, check if the dataset named ds
    has a child filesystem that is in the exclude list. If so, ds cannot be
    recursively snapshotted. """
    sds=ds.split('/')
    for exc in excludes:
        sexc=exc.split('/')
        if len(sds) > len(sexc):
            next
        if sds == sexc[:len(sds)]:
            return False
    return True

def narrow_recursive_filesystems(recursive_list):
    final_list=[]
    for ds in recursive_list:
        sds=ds.split('/')
        found=False
        logging.debug("checking if %s has a parent in %s" % (ds,recursive_list))
        for tmp in recursive_list:
            logging.debug("comparing %s to %s" % (tmp,ds))
            if tmp!=ds:
                stmp=tmp.split('/')
                if stmp == sds[:len(stmp)]:
                    found=True
                    next
        if found==False:
            final_list.append(ds)
    return final_list

def get_userprop_datasets(label="daily", userprop_name=USERPROP_NAME):
    """ This builds two lists of datasets - RECURSIVE_LIST and SINGLE_LIST
    based on the value of ZFS user properties com.sun:auto-snapshot and
    com.sun:auto-snapshot:${label}
    RECURSIVE_LIST is a list of datasets that can be snapshotted with -r
    SINGLE_LIST is a list of datasets to snapshot individually.
    """

    props=['name',
           userprop_name,
           SEP.join([userprop_name,label]) ]

    r=zfs_list(sort='name', properties=props)
    save=[]
    exclude=[]
    for row in r:
        save.append(row)
        if row[2] == 'false' or (row[1] == 'false' and row[2] == '-') or \
           (row[1] == '-' and row[2] == '-'):
            exclude.append(row[0])

    recursive_list=[]
    single_list=[]

    # find recursive snapshot sets
    for row in save:
        ds=row[0]
        logging.debug("checking %s" % ds)
        if can_recursive_snapshot(ds,exclude):
            logging.debug("OK to recursive snapshot %s" % ds)
            recursive_list.append(ds)
        elif ds not in exclude:
            logging.debug("OK to snapshot sole dataset %s" % ds)
            single_list.append(ds)
        else:
            logging.debug("NOT OK to snapshot %s" % ds)

    logging.debug("Pre-Narrowed recursive list is %s" % recursive_list)
    final_recursive_list=narrow_recursive_filesystems(recursive_list)
    logging.debug("Final recursive list is %s" % final_recursive_list)
    return (single_list,final_recursive_list)

<|endoftext|>