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

"""Load dependencies needed to compile and test the google-cloud-cpp library."""

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file")

def google_cloud_cpp_deps():
    """Loads dependencies need to compile the google-cloud-cpp library.

    Application developers can call this function from their WORKSPACE file
    to obtain all the necessary dependencies for google-cloud-cpp, including
    gRPC and its dependencies. This function only loads dependencies that
    have not been previously loaded, allowing application developers to
    override the version of the dependencies they want to use.
    """

    # Load rules_cc, used by googletest
    if "rules_cc" not in native.existing_rules():
        http_archive(
            name = "rules_cc",
            urls = [
                "https://github.com/bazelbuild/rules_cc/releases/download/0.0.1/rules_cc-0.0.1.tar.gz",
            ],
            sha256 = "4dccbfd22c0def164c8f47458bd50e0c7148f3d92002cdb459c2a96a68498241",
        )

    # Load Abseil
    if "com_google_absl" not in native.existing_rules():
        http_archive(
            name = "com_google_absl",
            strip_prefix = "abseil-cpp-20211102.0",
            urls = [
                "https://github.com/abseil/abseil-cpp/archive/20211102.0.tar.gz",
            ],
            sha256 = "dcf71b9cba8dc0ca9940c4b316a0c796be8fab42b070bb6b7cab62b48f0e66c4",
        )

    # Load a version of googletest that we know works.
    if "com_google_googletest" not in native.existing_rules():
        http_archive(
            name = "com_google_googletest",
            strip_prefix = "googletest-release-1.11.0",
            urls = [
                "https://github.com/google/googletest/archive/release-1.11.0.tar.gz",
            ],
            sha256 = "b4870bf121ff7795ba20d20bcdd8627b8e088f2d1dab299a031c1034eddc93d5",
        )

    # Load a version of benchmark that we know works.
    if "com_google_benchmark" not in native.existing_rules():
        http_archive(
            name = "com_google_benchmark",
            strip_prefix = "benchmark-1.6.0",
            urls = [
                "https://github.com/google/benchmark/archive/v1.6.0.tar.gz",
            ],
            sha256 = "1f71c72ce08d2c1310011ea6436b31e39ccab8c2db94186d26657d41747c85d6",
        )

    # Load the googleapis dependency.
    if "com_google_googleapis" not in native.existing_rules():
        http_archive(
            name = "com_google_googleapis",
            urls = [
                "https://github.com/googleapis/googleapis/archive/eddd3d845b5d462c0f2fb2d1bf764b921ae70a93.tar.gz",
            ],
            strip_prefix = "googleapis-eddd3d845b5d462c0f2fb2d1bf764b921ae70a93",
            sha256 = "adef434b41da6b7516fe5038bd6d8d5ea33b8031f6a3b2b4b6cdc5fc8ccb74dd",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:googleapis.BUILD",
            # Scaffolding for patching googleapis after download. For example:
            #   patches = ["googleapis.patch"]
            # NOTE: This should only be used while developing with a new
            # protobuf message. No changes to `patches` should ever be
            # committed to the main branch.
            patch_tool = "patch",
            patch_args = ["-p1"],
            patches = [],
        )

    # Load protobuf.
    if "com_google_protobuf" not in native.existing_rules():
        http_archive(
            name = "com_google_protobuf",
            strip_prefix = "protobuf-3.19.1",
            urls = [
                "https://github.com/protocolbuffers/protobuf/archive/v3.19.1.tar.gz",
            ],
            sha256 = "87407cd28e7a9c95d9f61a098a53cf031109d451a7763e7dd1253abf8b4df422",
        )

    # Load BoringSSL, this is used by gRPC, but as I write this (2021-06-03, circa gRPC-1.37.1), the version used by
    # gRPC does not compile with GCC-11
    if "boringssl" not in native.existing_rules():
        http_archive(
            name = "boringssl",
            # Use github mirror instead of https://boringssl.googlesource.com/boringssl
            # to obtain a boringssl archive with consistent sha256
            sha256 = "f0f433106f98f6f50ed6bbc2169f7c42dd73d13d0f09d431082519b5649903a6",
            strip_prefix = "boringssl-56eb68f64215738552be452e311da12047934ab4",
            urls = [
                "https://github.com/google/boringssl/archive/56eb68f64215738552be452e311da12047934ab4.tar.gz",
            ],
        )

    # Load gRPC and its dependencies, using a similar pattern to this function.
    if "com_github_grpc_grpc" not in native.existing_rules():
        http_archive(
            name = "com_github_grpc_grpc",
            strip_prefix = "grpc-1.42.0",
            urls = [
                "https://github.com/grpc/grpc/archive/v1.42.0.tar.gz",
            ],
            sha256 = "b2f2620c762427bfeeef96a68c1924319f384e877bc0e084487601e4cc6e434c",
        )

    # We use the cc_proto_library() rule from @com_google_protobuf, which
    # assumes that grpc_cpp_plugin and grpc_lib are in the //external: module
    native.bind(
        name = "grpc_cpp_plugin",
        actual = "@com_github_grpc_grpc//src/compiler:grpc_cpp_plugin",
    )

    native.bind(
        name = "grpc_lib",
        actual = "@com_github_grpc_grpc//:grpc++",
    )

    # We need libcurl for the Google Cloud Storage client.
    if "com_github_curl_curl" not in native.existing_rules():
        http_archive(
            name = "com_github_curl_curl",
            urls = [
                "https://curl.haxx.se/download/curl-7.69.1.tar.gz",
            ],
            strip_prefix = "curl-7.69.1",
            sha256 = "01ae0c123dee45b01bbaef94c0bc00ed2aec89cb2ee0fd598e0d302a6b5e0a98",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:curl.BUILD",
        )

    # We need the nlohmann_json library
    if "com_github_nlohmann_json" not in native.existing_rules():
        http_archive(
            name = "com_github_nlohmann_json",
            urls = [
                "https://github.com/nlohmann/json/releases/download/v3.10.4/include.zip",
            ],
            sha256 = "62c585468054e2d8e7c2759c0d990fd339d13be988577699366fe195162d16cb",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:nlohmann_json.BUILD",
        )

    # Load google/crc32c, a library to efficiently compute CRC32C checksums.
    if "com_github_google_crc32c" not in native.existing_rules():
        http_archive(
            name = "com_github_google_crc32c",
            strip_prefix = "crc32c-1.1.2",
            urls = [
                "https://github.com/google/crc32c/archive/1.1.2.tar.gz",
            ],
            sha256 = "ac07840513072b7fcebda6e821068aa04889018f24e10e46181068fb214d7e56",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:crc32c.BUILD",
        )
<commit_msg>chore(deps): update dependency com_github_grpc_grpc to v1.43.0 (#7748)<commit_after># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Load dependencies needed to compile and test the google-cloud-cpp library."""

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file")

def google_cloud_cpp_deps():
    """Loads dependencies need to compile the google-cloud-cpp library.

    Application developers can call this function from their WORKSPACE file
    to obtain all the necessary dependencies for google-cloud-cpp, including
    gRPC and its dependencies. This function only loads dependencies that
    have not been previously loaded, allowing application developers to
    override the version of the dependencies they want to use.
    """

    # Load rules_cc, used by googletest
    if "rules_cc" not in native.existing_rules():
        http_archive(
            name = "rules_cc",
            urls = [
                "https://github.com/bazelbuild/rules_cc/releases/download/0.0.1/rules_cc-0.0.1.tar.gz",
            ],
            sha256 = "4dccbfd22c0def164c8f47458bd50e0c7148f3d92002cdb459c2a96a68498241",
        )

    # Load Abseil
    if "com_google_absl" not in native.existing_rules():
        http_archive(
            name = "com_google_absl",
            strip_prefix = "abseil-cpp-20211102.0",
            urls = [
                "https://github.com/abseil/abseil-cpp/archive/20211102.0.tar.gz",
            ],
            sha256 = "dcf71b9cba8dc0ca9940c4b316a0c796be8fab42b070bb6b7cab62b48f0e66c4",
        )

    # Load a version of googletest that we know works.
    if "com_google_googletest" not in native.existing_rules():
        http_archive(
            name = "com_google_googletest",
            strip_prefix = "googletest-release-1.11.0",
            urls = [
                "https://github.com/google/googletest/archive/release-1.11.0.tar.gz",
            ],
            sha256 = "b4870bf121ff7795ba20d20bcdd8627b8e088f2d1dab299a031c1034eddc93d5",
        )

    # Load a version of benchmark that we know works.
    if "com_google_benchmark" not in native.existing_rules():
        http_archive(
            name = "com_google_benchmark",
            strip_prefix = "benchmark-1.6.0",
            urls = [
                "https://github.com/google/benchmark/archive/v1.6.0.tar.gz",
            ],
            sha256 = "1f71c72ce08d2c1310011ea6436b31e39ccab8c2db94186d26657d41747c85d6",
        )

    # Load the googleapis dependency.
    if "com_google_googleapis" not in native.existing_rules():
        http_archive(
            name = "com_google_googleapis",
            urls = [
                "https://github.com/googleapis/googleapis/archive/eddd3d845b5d462c0f2fb2d1bf764b921ae70a93.tar.gz",
            ],
            strip_prefix = "googleapis-eddd3d845b5d462c0f2fb2d1bf764b921ae70a93",
            sha256 = "adef434b41da6b7516fe5038bd6d8d5ea33b8031f6a3b2b4b6cdc5fc8ccb74dd",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:googleapis.BUILD",
            # Scaffolding for patching googleapis after download. For example:
            #   patches = ["googleapis.patch"]
            # NOTE: This should only be used while developing with a new
            # protobuf message. No changes to `patches` should ever be
            # committed to the main branch.
            patch_tool = "patch",
            patch_args = ["-p1"],
            patches = [],
        )

    # Load protobuf.
    if "com_google_protobuf" not in native.existing_rules():
        http_archive(
            name = "com_google_protobuf",
            strip_prefix = "protobuf-3.19.1",
            urls = [
                "https://github.com/protocolbuffers/protobuf/archive/v3.19.1.tar.gz",
            ],
            sha256 = "87407cd28e7a9c95d9f61a098a53cf031109d451a7763e7dd1253abf8b4df422",
        )

    # Load BoringSSL, this is used by gRPC, but as I write this (2021-06-03, circa gRPC-1.37.1), the version used by
    # gRPC does not compile with GCC-11
    if "boringssl" not in native.existing_rules():
        http_archive(
            name = "boringssl",
            # Use github mirror instead of https://boringssl.googlesource.com/boringssl
            # to obtain a boringssl archive with consistent sha256
            sha256 = "f0f433106f98f6f50ed6bbc2169f7c42dd73d13d0f09d431082519b5649903a6",
            strip_prefix = "boringssl-56eb68f64215738552be452e311da12047934ab4",
            urls = [
                "https://github.com/google/boringssl/archive/56eb68f64215738552be452e311da12047934ab4.tar.gz",
            ],
        )

    # Load gRPC and its dependencies, using a similar pattern to this function.
    if "com_github_grpc_grpc" not in native.existing_rules():
        http_archive(
            name = "com_github_grpc_grpc",
            strip_prefix = "grpc-1.43.0",
            urls = [
                "https://github.com/grpc/grpc/archive/v1.43.0.tar.gz",
            ],
            sha256 = "9647220c699cea4dafa92ec0917c25c7812be51a18143af047e20f3fb05adddc",
        )

    # We use the cc_proto_library() rule from @com_google_protobuf, which
    # assumes that grpc_cpp_plugin and grpc_lib are in the //external: module
    native.bind(
        name = "grpc_cpp_plugin",
        actual = "@com_github_grpc_grpc//src/compiler:grpc_cpp_plugin",
    )

    native.bind(
        name = "grpc_lib",
        actual = "@com_github_grpc_grpc//:grpc++",
    )

    # We need libcurl for the Google Cloud Storage client.
    if "com_github_curl_curl" not in native.existing_rules():
        http_archive(
            name = "com_github_curl_curl",
            urls = [
                "https://curl.haxx.se/download/curl-7.69.1.tar.gz",
            ],
            strip_prefix = "curl-7.69.1",
            sha256 = "01ae0c123dee45b01bbaef94c0bc00ed2aec89cb2ee0fd598e0d302a6b5e0a98",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:curl.BUILD",
        )

    # We need the nlohmann_json library
    if "com_github_nlohmann_json" not in native.existing_rules():
        http_archive(
            name = "com_github_nlohmann_json",
            urls = [
                "https://github.com/nlohmann/json/releases/download/v3.10.4/include.zip",
            ],
            sha256 = "62c585468054e2d8e7c2759c0d990fd339d13be988577699366fe195162d16cb",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:nlohmann_json.BUILD",
        )

    # Load google/crc32c, a library to efficiently compute CRC32C checksums.
    if "com_github_google_crc32c" not in native.existing_rules():
        http_archive(
            name = "com_github_google_crc32c",
            strip_prefix = "crc32c-1.1.2",
            urls = [
                "https://github.com/google/crc32c/archive/1.1.2.tar.gz",
            ],
            sha256 = "ac07840513072b7fcebda6e821068aa04889018f24e10e46181068fb214d7e56",
            build_file = "@com_github_googleapis_google_cloud_cpp//bazel:crc32c.BUILD",
        )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division
from acq4.modules.TaskRunner.analysisModules import AnalysisModule
from acq4.Manager import getManager
from PyQt4 import QtCore, QtGui
from imagingTemplate import Ui_Form
import numpy as np
import acq4.pyqtgraph as pg
import acq4.util.functions as fn
import acq4.util.metaarray as metaarray
from acq4.util.HelpfulException import HelpfulException
# import acq4.devices.Scanner.ScanUtilityFuncs as SUFA
from acq4.devices.Scanner.ScanProgram.rect import RectScan
from acq4.pyqtgraph.parametertree import ParameterTree, Parameter

class ImagingModule(AnalysisModule):
    def __init__(self, *args):
        AnalysisModule.__init__(self, *args)
        # self.ui = Ui_Form()
        # self.ui.setupUi(self)
        self.layout = QtGui.QGridLayout()
        self.setLayout(self.layout)
        self.splitter = QtGui.QSplitter()
        self.layout.addWidget(self.splitter)
        self.ptree = ParameterTree()
        self.splitter.addWidget(self.ptree)
        self.imageView = pg.ImageView()
        self.splitter.addWidget(self.imageView)

        # self.postGuiInit()

        self.params = Parameter(name='imager', children=[
            dict(name='scanner', type='interface', interfaceTypes=['scanner']),
            dict(name='detectors', type='group', addText="Add detector.."),
            dict(name='decomb', type='float', value=20e-6, suffix='s', siPrefix=True, bounds=[0, 1e-3], step=1e-6, children=[
                dict(name='auto', type='bool', value=True),
                dict(name='subpixel', type='bool', value=False),
                ]),
            dict(name='downsample', type='int', value=1, suffix='x', bounds=[1,None]),
            dict(name='display', type='bool', value=True),
            ])
        self.ptree.setParameters(self.params, showTop=False)
        self.params.sigTreeStateChanged.connect(self.update)
        self.params.child('detectors').sigAddNew.connect(self.addNewDetector)

        self.man = getManager()
        self.lastFrame = None
        # self.SUF = SUFA.ScannerUtilities()
        # self.ui.alphaSlider.valueChanged.connect(self.imageAlphaAdjust)
        self.img = pg.ImageItem()  ## image shown in camera module
        self.img.setLookupTable(self.imageView.ui.histogram.getLookupTable)  # image fetches LUT from the ImageView
        self.imageView.ui.histogram.sigLevelsChanged.connect(self._updateCamModImage)
        self.imageView.imageItem.setAutoDownsample(True)
        # self.ui.scannerComboBox.setTypes('scanner')
        # self.ui.detectorComboBox.setTypes('daqChannelGroup')

    def addNewDetector(self):
        self.params.child('detectors').addChild(
            dict(name='detector', type='interface', interfaceTypes=['daqChannelGroup']),
            autoIncrementName=True)

    def quit(self):
        self.clear()
        AnalysisModule.quit(self)

    def saveState(self):
        return self.params.saveState(filter='user')

    def restoreState(self, state):
        # for backward compat:
        det = state['children'].pop('detector', None)
        if det is not None:
            state['children']['detectors']['children']['detector'] = det

        self.params.restoreState(state, removeChildren=False)

    def taskSequenceStarted(self, *args):
        pass

    def taskFinished(self):
        pass

    def newFrame(self, frame):
        """
        Called when task is finished (truly completes, no errors/abort)
        frame contains all of the data returned from all devices
        """
        self.lastFrame = frame
        self.update()  # updates image

        # Store image if requested
        storeFlag = frame['cmd']['protocol']['storeData'] # get flag
        if storeFlag and len(self.lastResult) > 0:
            result = self.lastResult[0]  # for now we only handle single-component programs

            dirhandle = frame['cmd']['protocol']['storageDir'] # grab storage directory

            # to line below, add x, y for the camera (look at camera video output)

            info = [dict(name='Time', units='s', values=result['scanParams'].frameTimes()),
                    dict(name='X'), dict(name='Y'), {
                        'transform': result['transform'],
                        'imageProcessing': result['params'],
                    }]
            ma = metaarray.MetaArray(result['image'], info=info)
            fh = dirhandle.writeFile(ma, 'Imaging.ma')
            fh.setInfo(transform=result['transform'])

    def update(self):
        self.lastResult = []
        frame = self.lastFrame
        if frame is None:
            self.clear()
            return
        # imageDownSample = self.ui.downSampling.value() # this is the "image" downsample,
        # get the downsample for the daq. This is far more complicated than it should be...

        # Get PMT signal(s)
        pmtdata = []
        for detector in self.params.param('detectors'):
            data = frame['result'][detector.value()]["Channel":'Input']
            t = data.xvals('Time')
            pmtdata.append(data.asarray())

        if len(pmtdata) == 0:
            return

        # parse program options
        progs = frame['cmd'][self.params['scanner']]['program']
        if len(progs) == 0:
            self.image.setImage(np.zeros((1,1)))
            return

        # For now, we only support single-component scan programs.
        prog = progs[0]

        if prog['type'] == 'rect':
            # keep track of some analysis in case it should be stored later
            result = {'params': self.params.saveState(filter='user')['children']}
            self.lastResult.append(result)

            rs = RectScan()
            rs.restoreState(prog['scanInfo'])
            result['scanParams'] = rs

            # Determine decomb duration
            auto = self.params['decomb', 'auto']
            if auto:
                (decombed, lag) = rs.measureMirrorLag(pmtdata[0])
                self.params['decomb'] = lag
            decomb = self.params['decomb']

            # Extract from PMT array
            imageData = []
            for chan in pmtdata:
                chanImage = rs.extractImage(chan, offset=decomb, subpixel=self.params['decomb', 'subpixel'])
                imageData.append(chanImage.reshape(chanImage.shape + (1,)))

            if len(imageData) == 1:
                imageData = imageData[0]
                levelMode = 'mono'
            else:
                if len(imageData) == 2:
                    imageData.append(np.zeros(imageData[0].shape, dtype=imageData[0].dtype))
                imageData = np.concatenate(imageData, axis=-1)
                levelMode = 'rgba'

            if imageData.size == 0:
                self.clear()
                raise Exception('image Data has zero size')

            # Downsample
            ds = self.params['downsample']
            if ds > 1:
                imageData = pg.downsample(imageData, ds, axis=2)

            # Collected as (frame, row, col) but pg prefers images like (frame, col, row)
            imageData = imageData.transpose((0, 2, 1, 3)[:imageData.ndim])
            result['image'] = imageData

            # compute global transform
            tr = rs.imageTransform()
            st = pg.QtGui.QTransform()
            st.scale(self.params['downsample'], 1)
            tr = st * tr
            result['transform'] = pg.SRTTransform3D(tr)

            # Display image locally
            self.imageView.setImage(imageData, levelMode=levelMode)
            self.imageView.getView().setAspectLocked(True)
#            self.imageView.imageItem.setRect(QtCore.QRectF(0., 0., rs.width, rs.height))  # TODO: rs.width and rs.height might not be correct!
            self.imageView.imageItem.resetTransform()
            self.imageView.imageItem.scale((rs.width/rs.height)/(imageData.shape[1]/imageData.shape[2]), 1.0)
            self.imageView.autoRange()

            # Display image remotely (in the same camera module as used by the scanner device)
            if self.params['display']:
                self.img.setVisible(True)
                sd = self.pr.getDevice(self.params['scanner'])
                camMod = sd.cameraModule().window()
                camMod.addItem(self.img)
                self.img.setImage(imageData.mean(axis=0))
                self.img.setTransform(tr)
            else:
                self.img.setVisible(False)

        if prog['type'] == 'multipleLineScan':
            totSamps = int(np.sum(prog['scanPointList'])) # samples per scan, before downsampling
            imageData = pmtdata[prog['startStopIndices'][0]:prog['startStopIndices'][0]+int((nscans*totSamps))].copy()
            imageData = imageData.reshape(nscans, totSamps)
            csum = np.cumsum(prog['scanPointList'])
            for i in xrange(0, len(csum), 2):
                    imageData[:,csum[i]:csum[i+1]] = 0

            #imageData = imageData[:,0:prog['samplesPerScan']] # trim off the pause data
            imageData = fn.downsample(imageData, imageDownSample, axis=1)
            self.ui.plotWidget.setImage(imageData)
            self.ui.plotWidget.getView().setAspectLocked(False)
            self.ui.plotWidget.imageItem.setRect(QtCore.QRectF(startT, 0.0, totSamps*dt*nscans, dist ))
            self.ui.plotWidget.autoRange()
            storeFlag = frame['cmd']['protocol']['storeData'] # get flag
            if storeFlag:
                dirhandle = frame['cmd']['protocol']['storageDir'] # grab directory
                self.info={'detector': self.detectorDevice(), 'scanner': self.scannerDevice(), 'indices': prog['startStopIndices'],
                           'scanPointList': prog['scanPointList'], 'nscans': prog['nScans'],
                           'positions': prog['points'],
                           'downSample': imageDownSample, 'daqDownSample': daqDownSample}
                info = [dict(name='Time', units='s',
                             values=t[prog['startStopIndices'][0]:prog['startStopIndices'][1]-int(totSamps):int(totSamps)]),
                        dict(name='Distance'), self.info]
                ma = metaarray.MetaArray(imageData, info=info)
                dirhandle.writeFile(ma, 'Imaging.ma')

        if prog['type'] == 'rectScan':
            samplesPerScan = prog['imageSize'][0]*prog['imageSize'][1]
            getManager().data = prog, pmtdata
            imageData = pmtdata[prog['startStopIndices'][0]:prog['startStopIndices'][0] + nscans*samplesPerScan]
            imageData=imageData.reshape(nscans, prog['imageSize'][1], prog['imageSize'][0])
            imageData = imageData.transpose(0,2,1)
            imageAve = np.mean(imageData, axis=0)
            #print prog['scanParameters']
            self.SUF.setScannerParameters(prog['scanParameters']) # load up information for Scanner calculations
            self.SUF.setScanInfo(prog['scanInfo'])
            imageAve, bestShift = self.SUF.adjustBidirectional(imageAve, True, 0.)
            #print 'bestShift: ', bestShift, ' microseconds'
            bestShift = 200e-6
            for i in range(nscans):
                (decombedImage, shift) = self.SUF.adjustBidirectional(imageData[i], False, bestShift)
                roImage = self.SUF.removeOverscan(decombedImage)
                if i == 0:
                    newImage = np.zeros((nscans, roImage.shape[0], roImage.shape[1]))
                newImage[i] = roImage
            #print imageAve.shape
            #print newImage.shape
            imageData = newImage

            # imageData = fn.downsample(imageData, imageDownSample, axis=0)
            if imageData.size == 0:
                raise Exception('image Data has zero size')
            self.ui.plotWidget.setImage(imageData)
            pts = prog['points']
            floatpoints =[ (float(x[0]), float(x[1])) for x in pts]
            width  = (pts[1] -pts[0]).length() # width is x in M
            height = (pts[2]- pts[0]).length() # heigh in M
            self.ui.plotWidget.getView().setAspectLocked(True)
            self.ui.plotWidget.imageItem.setRect(QtCore.QRectF(0., 0., width, height))
            self.ui.plotWidget.autoRange()
           # self.ui.histogram.imageChanged(autoLevel=True)
            #print 'rectscan - nscans, samplesperscan: ', prog['nScans'], samplesPerScan
            sd = self.pr.getDevice(self.scannerDevice())
            camMod = sd.cameraModule().window()
            if self.img is not None:
                camMod.removeItem(self.img)
                self.img = None
            self.img = pg.ImageItem(imageData.mean(axis=0))
            camMod.addItem(self.img)
            w = imageData.shape[1]
            h = imageData.shape[2]
            localPts = map(pg.Vector, [[0,0], [w,0], [0,h], [0,0,1]]) # w and h of data of image in pixels.
            globalPts = prog['points'] # sort of. -
            m = pg.solve3DTransform(localPts, map(pg.Vector, globalPts+[[0,0,1]]))
            m[:,2] = m[:,3]
            m[2] = m[3]
            m[2,2] = 1
            tr = QtGui.QTransform(*m[:3,:3].transpose().reshape(9))
            self.img.setTransform(tr)
            storeFlag = frame['cmd']['protocol']['storeData'] # get flag
           # print 'srttransform: ', pg.SRTTransform3D(tr)
            if storeFlag:
                dirhandle = frame['cmd']['protocol']['storageDir'] # grab directory
                self.info={'detector': self.detectorDevice(), 'scanner': self.scannerDevice(), 'indices': prog['startStopIndices'],
                           'samplesPerScan': samplesPerScan, 'nscans': prog['nScans'],
                           'positions': floatpoints, # prog['points'],
                           'downSample': imageDownSample,
                           'transform': pg.SRTTransform3D(tr),
                           }

                # to line below, add x, y for the camera (look at camera video output)
                info = [dict(name='Time', units='s',
                             values=t[prog['startStopIndices'][0]:prog['startStopIndices'][0]+nscans*samplesPerScan:samplesPerScan]),
                        dict(name='Distance'), self.info]
                print self.info
                print info
                ma = metaarray.MetaArray(imageData, info=info)

                dirhandle.writeFile(ma, 'Imaging.ma')

    def clear(self):
        self.imageView.clear()
        scene = self.img.scene()
        if scene is not None:
            scene.removeItem(self.img)

    def imageAlphaAdjust(self):
        if self.img is None:
            return
        alpha = self.ui.alphaSlider.value()
        self.img.setImage(opacity=float(alpha/100.))

    def _updateCamModImage(self):
        # update image levels
        self.img.setLevels(self.imageView.ui.histogram.getLevels())
<commit_msg>fix Imaging mod. save/restore of detectors<commit_after># -*- coding: utf-8 -*-
from __future__ import division
from acq4.modules.TaskRunner.analysisModules import AnalysisModule
from acq4.Manager import getManager
from PyQt4 import QtCore, QtGui
from imagingTemplate import Ui_Form
import numpy as np
import acq4.pyqtgraph as pg
import acq4.util.functions as fn
import acq4.util.metaarray as metaarray
from acq4.util.HelpfulException import HelpfulException
# import acq4.devices.Scanner.ScanUtilityFuncs as SUFA
from acq4.devices.Scanner.ScanProgram.rect import RectScan
from acq4.pyqtgraph.parametertree import ParameterTree, Parameter

class ImagingModule(AnalysisModule):
    def __init__(self, *args):
        AnalysisModule.__init__(self, *args)
        # self.ui = Ui_Form()
        # self.ui.setupUi(self)
        self.layout = QtGui.QGridLayout()
        self.setLayout(self.layout)
        self.splitter = QtGui.QSplitter()
        self.layout.addWidget(self.splitter)
        self.ptree = ParameterTree()
        self.splitter.addWidget(self.ptree)
        self.imageView = pg.ImageView()
        self.splitter.addWidget(self.imageView)

        # self.postGuiInit()

        self.params = Parameter(name='imager', children=[
            dict(name='scanner', type='interface', interfaceTypes=['scanner']),
            dict(name='detectors', type='group', addText="Add detector.."),
            dict(name='decomb', type='float', value=20e-6, suffix='s', siPrefix=True, bounds=[0, 1e-3], step=1e-6, children=[
                dict(name='auto', type='bool', value=True),
                dict(name='subpixel', type='bool', value=False),
                ]),
            dict(name='downsample', type='int', value=1, suffix='x', bounds=[1,None]),
            dict(name='display', type='bool', value=True),
            ])
        self.ptree.setParameters(self.params, showTop=False)
        self.params.sigTreeStateChanged.connect(self.update)
        self.params.child('detectors').sigAddNew.connect(self.addNewDetector)

        self.man = getManager()
        self.lastFrame = None
        # self.SUF = SUFA.ScannerUtilities()
        # self.ui.alphaSlider.valueChanged.connect(self.imageAlphaAdjust)
        self.img = pg.ImageItem()  ## image shown in camera module
        self.img.setLookupTable(self.imageView.ui.histogram.getLookupTable)  # image fetches LUT from the ImageView
        self.imageView.ui.histogram.sigLevelsChanged.connect(self._updateCamModImage)
        self.imageView.imageItem.setAutoDownsample(True)
        # self.ui.scannerComboBox.setTypes('scanner')
        # self.ui.detectorComboBox.setTypes('daqChannelGroup')

    def addNewDetector(self, name='detector', value=None):
        self.params.child('detectors').addChild(
            dict(name=name, type='interface', interfaceTypes=['daqChannelGroup'], value=value),
            autoIncrementName=True)

    def quit(self):
        self.clear()
        AnalysisModule.quit(self)

    def saveState(self):
        return self.params.saveState(filter='user')

    def restoreState(self, state):
        detectors = {}

        # for backward compat:
        det = state['children'].pop('detector', None)
        if det is not None:
            detectors['detector'] = det['value']
        # current format:
        dets = state['children'].pop('detectors', {})
        for name, data in dets['children'].items():
            detectors[name] = data['value']

        self.params.restoreState(state, removeChildren=False)
        for name, det in detectors.items():
            self.addNewDetector(name, det)

    def taskSequenceStarted(self, *args):
        pass

    def taskFinished(self):
        pass

    def newFrame(self, frame):
        """
        Called when task is finished (truly completes, no errors/abort)
        frame contains all of the data returned from all devices
        """
        self.lastFrame = frame
        self.update()  # updates image

        # Store image if requested
        storeFlag = frame['cmd']['protocol']['storeData'] # get flag
        if storeFlag and len(self.lastResult) > 0:
            result = self.lastResult[0]  # for now we only handle single-component programs

            dirhandle = frame['cmd']['protocol']['storageDir'] # grab storage directory

            # to line below, add x, y for the camera (look at camera video output)

            info = [dict(name='Time', units='s', values=result['scanParams'].frameTimes()),
                    dict(name='X'), dict(name='Y'), {
                        'transform': result['transform'],
                        'imageProcessing': result['params'],
                    }]
            ma = metaarray.MetaArray(result['image'], info=info)
            fh = dirhandle.writeFile(ma, 'Imaging.ma')
            fh.setInfo(transform=result['transform'])

    def update(self):
        self.lastResult = []
        frame = self.lastFrame
        if frame is None:
            self.clear()
            return
        # imageDownSample = self.ui.downSampling.value() # this is the "image" downsample,
        # get the downsample for the daq. This is far more complicated than it should be...

        # Get PMT signal(s)
        pmtdata = []
        for detector in self.params.param('detectors'):
            data = frame['result'][detector.value()]["Channel":'Input']
            t = data.xvals('Time')
            pmtdata.append(data.asarray())

        if len(pmtdata) == 0:
            return

        # parse program options
        progs = frame['cmd'][self.params['scanner']]['program']
        if len(progs) == 0:
            self.image.setImage(np.zeros((1,1)))
            return

        # For now, we only support single-component scan programs.
        prog = progs[0]

        if prog['type'] == 'rect':
            # keep track of some analysis in case it should be stored later
            result = {'params': self.params.saveState(filter='user')['children']}
            self.lastResult.append(result)

            rs = RectScan()
            rs.restoreState(prog['scanInfo'])
            result['scanParams'] = rs

            # Determine decomb duration
            auto = self.params['decomb', 'auto']
            if auto:
                (decombed, lag) = rs.measureMirrorLag(pmtdata[0])
                self.params['decomb'] = lag
            decomb = self.params['decomb']

            # Extract from PMT array
            imageData = []
            for chan in pmtdata:
                chanImage = rs.extractImage(chan, offset=decomb, subpixel=self.params['decomb', 'subpixel'])
                imageData.append(chanImage.reshape(chanImage.shape + (1,)))

            if len(imageData) == 1:
                imageData = imageData[0]
                levelMode = 'mono'
            else:
                if len(imageData) == 2:
                    imageData.append(np.zeros(imageData[0].shape, dtype=imageData[0].dtype))
                imageData = np.concatenate(imageData, axis=-1)
                levelMode = 'rgba'

            if imageData.size == 0:
                self.clear()
                raise Exception('image Data has zero size')

            # Downsample
            ds = self.params['downsample']
            if ds > 1:
                imageData = pg.downsample(imageData, ds, axis=2)

            # Collected as (frame, row, col) but pg prefers images like (frame, col, row)
            imageData = imageData.transpose((0, 2, 1, 3)[:imageData.ndim])
            result['image'] = imageData

            # compute global transform
            tr = rs.imageTransform()
            st = pg.QtGui.QTransform()
            st.scale(self.params['downsample'], 1)
            tr = st * tr
            result['transform'] = pg.SRTTransform3D(tr)

            # Display image locally
            self.imageView.setImage(imageData, levelMode=levelMode)
            self.imageView.getView().setAspectLocked(True)
#            self.imageView.imageItem.setRect(QtCore.QRectF(0., 0., rs.width, rs.height))  # TODO: rs.width and rs.height might not be correct!
            self.imageView.imageItem.resetTransform()
            self.imageView.imageItem.scale((rs.width/rs.height)/(imageData.shape[1]/imageData.shape[2]), 1.0)
            self.imageView.autoRange()

            # Display image remotely (in the same camera module as used by the scanner device)
            if self.params['display']:
                self.img.setVisible(True)
                sd = self.pr.getDevice(self.params['scanner'])
                camMod = sd.cameraModule().window()
                camMod.addItem(self.img)
                self.img.setImage(imageData.mean(axis=0))
                self.img.setTransform(tr)
            else:
                self.img.setVisible(False)

        if prog['type'] == 'multipleLineScan':
            totSamps = int(np.sum(prog['scanPointList'])) # samples per scan, before downsampling
            imageData = pmtdata[prog['startStopIndices'][0]:prog['startStopIndices'][0]+int((nscans*totSamps))].copy()
            imageData = imageData.reshape(nscans, totSamps)
            csum = np.cumsum(prog['scanPointList'])
            for i in xrange(0, len(csum), 2):
                    imageData[:,csum[i]:csum[i+1]] = 0

            #imageData = imageData[:,0:prog['samplesPerScan']] # trim off the pause data
            imageData = fn.downsample(imageData, imageDownSample, axis=1)
            self.ui.plotWidget.setImage(imageData)
            self.ui.plotWidget.getView().setAspectLocked(False)
            self.ui.plotWidget.imageItem.setRect(QtCore.QRectF(startT, 0.0, totSamps*dt*nscans, dist ))
            self.ui.plotWidget.autoRange()
            storeFlag = frame['cmd']['protocol']['storeData'] # get flag
            if storeFlag:
                dirhandle = frame['cmd']['protocol']['storageDir'] # grab directory
                self.info={'detector': self.detectorDevice(), 'scanner': self.scannerDevice(), 'indices': prog['startStopIndices'],
                           'scanPointList': prog['scanPointList'], 'nscans': prog['nScans'],
                           'positions': prog['points'],
                           'downSample': imageDownSample, 'daqDownSample': daqDownSample}
                info = [dict(name='Time', units='s',
                             values=t[prog['startStopIndices'][0]:prog['startStopIndices'][1]-int(totSamps):int(totSamps)]),
                        dict(name='Distance'), self.info]
                ma = metaarray.MetaArray(imageData, info=info)
                dirhandle.writeFile(ma, 'Imaging.ma')

        if prog['type'] == 'rectScan':
            samplesPerScan = prog['imageSize'][0]*prog['imageSize'][1]
            getManager().data = prog, pmtdata
            imageData = pmtdata[prog['startStopIndices'][0]:prog['startStopIndices'][0] + nscans*samplesPerScan]
            imageData=imageData.reshape(nscans, prog['imageSize'][1], prog['imageSize'][0])
            imageData = imageData.transpose(0,2,1)
            imageAve = np.mean(imageData, axis=0)
            #print prog['scanParameters']
            self.SUF.setScannerParameters(prog['scanParameters']) # load up information for Scanner calculations
            self.SUF.setScanInfo(prog['scanInfo'])
            imageAve, bestShift = self.SUF.adjustBidirectional(imageAve, True, 0.)
            #print 'bestShift: ', bestShift, ' microseconds'
            bestShift = 200e-6
            for i in range(nscans):
                (decombedImage, shift) = self.SUF.adjustBidirectional(imageData[i], False, bestShift)
                roImage = self.SUF.removeOverscan(decombedImage)
                if i == 0:
                    newImage = np.zeros((nscans, roImage.shape[0], roImage.shape[1]))
                newImage[i] = roImage
            #print imageAve.shape
            #print newImage.shape
            imageData = newImage

            # imageData = fn.downsample(imageData, imageDownSample, axis=0)
            if imageData.size == 0:
                raise Exception('image Data has zero size')
            self.ui.plotWidget.setImage(imageData)
            pts = prog['points']
            floatpoints =[ (float(x[0]), float(x[1])) for x in pts]
            width  = (pts[1] -pts[0]).length() # width is x in M
            height = (pts[2]- pts[0]).length() # heigh in M
            self.ui.plotWidget.getView().setAspectLocked(True)
            self.ui.plotWidget.imageItem.setRect(QtCore.QRectF(0., 0., width, height))
            self.ui.plotWidget.autoRange()
           # self.ui.histogram.imageChanged(autoLevel=True)
            #print 'rectscan - nscans, samplesperscan: ', prog['nScans'], samplesPerScan
            sd = self.pr.getDevice(self.scannerDevice())
            camMod = sd.cameraModule().window()
            if self.img is not None:
                camMod.removeItem(self.img)
                self.img = None
            self.img = pg.ImageItem(imageData.mean(axis=0))
            camMod.addItem(self.img)
            w = imageData.shape[1]
            h = imageData.shape[2]
            localPts = map(pg.Vector, [[0,0], [w,0], [0,h], [0,0,1]]) # w and h of data of image in pixels.
            globalPts = prog['points'] # sort of. -
            m = pg.solve3DTransform(localPts, map(pg.Vector, globalPts+[[0,0,1]]))
            m[:,2] = m[:,3]
            m[2] = m[3]
            m[2,2] = 1
            tr = QtGui.QTransform(*m[:3,:3].transpose().reshape(9))
            self.img.setTransform(tr)
            storeFlag = frame['cmd']['protocol']['storeData'] # get flag
           # print 'srttransform: ', pg.SRTTransform3D(tr)
            if storeFlag:
                dirhandle = frame['cmd']['protocol']['storageDir'] # grab directory
                self.info={'detector': self.detectorDevice(), 'scanner': self.scannerDevice(), 'indices': prog['startStopIndices'],
                           'samplesPerScan': samplesPerScan, 'nscans': prog['nScans'],
                           'positions': floatpoints, # prog['points'],
                           'downSample': imageDownSample,
                           'transform': pg.SRTTransform3D(tr),
                           }

                # to line below, add x, y for the camera (look at camera video output)
                info = [dict(name='Time', units='s',
                             values=t[prog['startStopIndices'][0]:prog['startStopIndices'][0]+nscans*samplesPerScan:samplesPerScan]),
                        dict(name='Distance'), self.info]
                print self.info
                print info
                ma = metaarray.MetaArray(imageData, info=info)

                dirhandle.writeFile(ma, 'Imaging.ma')

    def clear(self):
        self.imageView.clear()
        scene = self.img.scene()
        if scene is not None:
            scene.removeItem(self.img)

    def imageAlphaAdjust(self):
        if self.img is None:
            return
        alpha = self.ui.alphaSlider.value()
        self.img.setImage(opacity=float(alpha/100.))

    def _updateCamModImage(self):
        # update image levels
        self.img.setLevels(self.imageView.ui.histogram.getLevels())
<|endoftext|>


<commit_before>"""based on beeswithmachineguns

* separated configuration and core functionality

* output is logged instead of printed

* class based approach
  (two reasons:
    * ease testing
    * functionality easily adjustable by by overwriting methods

* thrown out subnet stuff, as I don't understand the currrent implementation
  (seems somewhat nonsensical to me ...)
"""
import inspect
import simplejson as json
import logging
from plumbum.path import LocalPath, RemotePath, LocalWorkdir

from beeswithmachineguns.bees import *
from beeswithmachineguns.obj_inspection import obj_attr

log = logging.getLogger(__name__)

class BeeSting(Exception):
    """class for exceptions raised by this module"""
    pass

class JsonConfigger(object):
    def __init__(self, configPath):
        self._configPath = LocalPath(configPath)

    def load_config(self):
        config = json.loads(self._configPath.read())
        for key, value in config.items():
            log.debug("read from _config: %s <- %s", key, value)
            setattr(self, key, value)

    def save_config(self):
        self._configPath.write(json.dumps(self._config))

    def remove_config(self):
        self._configPath.delete()

    def get_config(self):
        return self._config

    @property
    def _config(self):
        attrs = {k: self._nrmlz(v) for k, v in self.__dict__.items()
                 if not k.startswith('_')}
        props = {k: self._nrmlz(getattr(self, k)) for k in self._publicProps}
        attrs.update(props.items())
        return attrs

    @property
    def _publicProps(self):
        return [name for (name, member) in inspect.getmembers(self.__class__)
                if not name.startswith('_') and type(member) == property]

    def _nrmlz(self, value):
        """normalize strange things to json compatible values"""
        if isinstance(value, str):
            return value

        if isinstance(value, LocalPath):
            return str(value)

        if isinstance(value, list):
            return [self._nrmlz(v) for v in value]

        return value

class ProjectConfig(JsonConfigger):
    CONFIG = 'hive.json'
    """Global configuration"""
    KEY_NAME_PREFIX = "aws-ec2"
    DEFAULT_SECURITY_GROUP = 'default'
    DEFAULT_ZONE = 'us-east-1d'
    DEFAULT_INSTANCE_ID = 'ami-ff17fb96'
    DEFAULT_NUMBER_OF_BEES = 10

    def __init__(self):
        super(ProjectConfig, self).__init__(self.CONFIG)
        self._keyContainerPaths = [self._configPath.dirname,
                                   LocalPath(os.getenv('HOME')) / '.ssh']
        self.securityGroupNames = [self.DEFAULT_SECURITY_GROUP]
        self.zone = self.DEFAULT_ZONE
        self.instanceId = self.DEFAULT_INSTANCE_ID
        self.numberOfBees = self.DEFAULT_NUMBER_OF_BEES

    @property
    def region(self):
        """calculate region from configured zone:
            * keep unchanged for gov zones
            * chop off the last letter in e.g. "us-east-1d"
        """
        return self.zone if 'gov' in self.zone else self.zone[:-1]

    @property
    def keyPath(self):
        if not hasattr(self, '_keyPath'):
            for path in self._keyContainerPaths:
                potentialKeyPath = path / self.keyName
                log.debug("try key path %s", potentialKeyPath)
                if potentialKeyPath.exists():
                    self._keyPath = potentialKeyPath
                    break
            else:
                raise BeeSting(
                    'no key named %s found in %s' %
                    (self.keyName, self._keyContainerPaths))

        return self._keyPath

    @property
    def keyName(self):
        return "%s-%s.pem" % (self.KEY_NAME_PREFIX, self.region)

class Ec2Connection(object):
    def __init__(self, **kwargs):
        self.zone = kwargs.get('zone')
        self.region = kwargs.get('region')
        self.securityGroupNames = kwargs.get('securityGroupNames')
        self.instanceId = kwargs.get('instanceId')
        self.numberOfBees = kwargs.get('numberOfBees')
        self.keyPath = kwargs.get('keyPath')
        self._connection = None

    @property
    def securityGroupIds(self):
        return [group.vpc_id for group in self.securityGroups]

    @property
    def securityGroups(self):
        return [g for g in self._allSecurityGroups
                if g.name in self.securityGroupNames]

    @property
    def _allSecurityGroups(self):
        return self.connection.get_all_security_groups()

    @property
    def connection(self):
        if not self._connection:
            self._connection = boto.ec2.connect_to_region(self.region)
        return self._connection

class CurrentHive(JsonConfigger):
    CONFIG = 'current.json'
    """configuration of an active bee hive (autogenerated)"""

    def __init__(self):
        super(CurrentHive, self).__init__(self.CONFIG)
        self.username = ''
        self.keyName = ''
        self.zone = ''
        self.beesIds = ''

    @property
    def isConfigured(self):
        return all(self._config)

    @property
    def isActive(self):
        return self._configPath.exists()

class Beekeeper(object):
    DEFAULT_INSTANCE_TYPE = 't1.micro'

    def __init__(self):
        # self.group = None
        # self.imageId = None
        self.instanceType = self.DEFAULT_INSTANCE_TYPE

    @property
    def hive(self):
        if not hasattr(self, '_hive'):
            self._hive = CurrentHive()
        return self._hive

    def up(self):
        if self.hive.isActive:
            log.warning("hive is up already: %s", self.hive.get_config())
            return

        pem_path = get_pem_path(self.hive.keyName)

        if not os.path.isfile(pem_path):
            print 'Warning. No key file found for %s. You will need to add this key to your SSH agent to connect.' % pem_path

        print 'Connecting to the hive.'

        ec2_connection = boto.ec2.connect_to_region(get_region(zone))

        print 'Attempting to call up %i bees.' % count

        reservation = ec2_connection.run_instances(
            image_id=image_id,
            min_count=count,
            max_count=count,
            key_name=self.keyName,
            security_groups=self.securityGroupIds,
            instance_type=instance_type,
            placement=None if 'gov' in zone else zone,
            subnet_id='')

        print 'Waiting for bees to load their machine guns...'

        instance_ids = []

        for instance in reservation.instances:
            instance.update()
            while instance.state != 'running':
                print '.'
                time.sleep(5)
                instance.update()

            instance_ids.append(instance.id)

            print 'Bee %s is ready for the attack.' % instance.id

        ec2_connection.create_tags(instance_ids, { "Name": "a bee!" })

        write_server_list(username, key_name, zone, reservation.instances)

        print 'The swarm has assembled %i bees.' % len(reservation.instances)

class LoggingConfig(object):
    NAME = 'bees'

    def __init__(self):
        self.workPath = LocalWorkdir()
        self.localLogPath = None
        """:type: LocalPath"""

    def init_logging(self, logLevel=logging.INFO, logToFile=True):
        log.setLevel(logLevel)
        self.localLogPath = self.workPath / (self.NAME + '.log')
        fmt = '%(asctime)s %(name)s %(levelname)s: %(message)s'
        logging.basicConfig(format=fmt)
        if logToFile:
            fh = logging.FileHandler(filename=str(self.localLogPath))
            fh.setFormatter(logging.Formatter(fmt))
            log.addHandler(fh)
        log.name = self.NAME if log.name == '__main__' else log.name
        log.debug("working in %s", self.workPath)

def main(workingPath=None):
    workPath = LocalWorkdir()
    if workingPath:
        workPath.chdir(workingPath)
    log.info('working in %s', workPath)
    logCnf = LoggingConfig()
    logCnf.init_logging(logLevel=logging.DEBUG)
    cnf = ProjectConfig()

    print cnf.numberOfBees
    print cnf._config

    ec2 = Ec2Connection(**cnf._config)
    # for g in ec2._allSecurityGroups:
    #     print obj_attr(g)
    print ec2.securityGroupIds
    exit()

    print ec2.connection
    print ec2.__dict__.items()
    print type(ec2.connection)

    exit()
    hive = CurrentHive()
    print hive.isActive
    if hive.isActive:
        hive.load_config()
    print hive.isActive
    # print hive._config
    # print hive.isConfigured
    print hive._config

if __name__ == '__main__':
    main('../tests/fake_project_dir')

    # time.sleep(0.01)
    # for k, v in cnf.__dict__.items():
    #     print k, v
    # print cnf.beesIds
<commit_msg>implemented up, but not tested yet<commit_after>"""based on beeswithmachineguns

* separated configuration and core functionality

* output is logged instead of printed

* class based approach
  (two reasons:
    * ease testing
    * functionality easily adjustable by by overwriting methods

* thrown out subnet stuff, as I don't understand the currrent implementation
  (seems somewhat nonsensical to me ...)
"""
import inspect
import simplejson as json
import logging
from plumbum.path import LocalPath, RemotePath, LocalWorkdir

from beeswithmachineguns.bees import *
from beeswithmachineguns.obj_inspection import obj_attr

log = logging.getLogger(__name__)

class BeeSting(Exception):
    """class for exceptions raised by this module"""
    pass

class JsonConfigger(object):
    def __init__(self, configPath):
        self._configPath = LocalPath(configPath)

    def load_config(self):
        config = json.loads(self._configPath.read())
        for key, value in config.items():
            log.debug("read from _config: %s <- %s", key, value)
            setattr(self, key, value)

    def save_config(self):
        self._configPath.write(json.dumps(self.get_config()))

    def remove_config(self):
        self._configPath.delete()

    def get_config(self):
        attrs = {k: self._nrmlz(v) for k, v in self.__dict__.items()
                 if not k.startswith('_')}
        props = {k: self._nrmlz(getattr(self, k)) for k in self._publicProps}
        attrs.update(props.items())
        return attrs

    @property
    def _publicProps(self):
        return [name for (name, member) in inspect.getmembers(self.__class__)
                if not name.startswith('_') and type(member) == property]

    def _nrmlz(self, value):
        """normalize strange things to json compatible values"""
        if isinstance(value, str):
            return value

        if isinstance(value, LocalPath):
            return str(value)

        if isinstance(value, list):
            return [self._nrmlz(v) for v in value]

        return value

class ProjectConfig(JsonConfigger):
    CONFIG = 'hive.json'
    """Global configuration"""
    KEY_NAME_PREFIX = "aws-ec2"
    DEFAULT_SECURITY_GROUP = 'default'
    DEFAULT_ZONE = 'us-east-1d'
    DEFAULT_INSTANCE_ID = 'ami-ff17fb96'
    DEFAULT_INSTANCE_TYPE = 't1.micro'

    DEFAULT_NUMBER_OF_BEES = 10

    def __init__(self):
        super(ProjectConfig, self).__init__(self.CONFIG)
        self._keyContainerPaths = [self._configPath.dirname,
                                   LocalPath(os.getenv('HOME')) / '.ssh']
        self.securityGroupNames = [self.DEFAULT_SECURITY_GROUP]
        self.zone = self.DEFAULT_ZONE
        self.instanceId = self.DEFAULT_INSTANCE_ID
        self.numberOfBees = self.DEFAULT_NUMBER_OF_BEES
        self.instanceType = self.DEFAULT_INSTANCE_TYPE

    @property
    def region(self):
        """calculate region from configured zone:
            * keep unchanged for gov zones
            * chop off the last letter in e.g. "us-east-1d"
        """
        return self.zone if self.isGovZone else self.zone[:-1]

    @property
    def isGovZone(self):
        return 'gov' in self.zone

    @property
    def keyPath(self):
        if not hasattr(self, '_keyPath'):
            for path in self._keyContainerPaths:
                potentialKeyPath = path / self.keyName
                log.debug("try key path %s", potentialKeyPath)
                if potentialKeyPath.exists():
                    self._keyPath = potentialKeyPath
                    break
            else:
                raise BeeSting(
                    'no key named %s found in %s' %
                    (self.keyName, self._keyContainerPaths))

        return self._keyPath

    @property
    def keyName(self):
        return "%s-%s.pem" % (self.KEY_NAME_PREFIX, self.region)

class Ec2Connection(object):
    def __init__(self, **kwargs):
        self.zone = kwargs.get('zone')
        self.isGovZone = kwargs.get('isGovZone')
        self.region = kwargs.get('region')
        self.securityGroupNames = kwargs.get('securityGroupNames')
        self.instanceId = kwargs.get('instanceId')
        self.instanceType = kwargs.get('instanceType')
        self.numberOfBees = kwargs.get('numberOfBees')
        self.keyName = LocalPath(kwargs.get('keyName'))
        self.keyPath = LocalPath(kwargs.get('keyPath'))
        self._connection = None

    @property
    def securityGroupIds(self):
        return [group.vpc_id for group in self.securityGroups]

    @property
    def securityGroups(self):
        return [g for g in self._allSecurityGroups
                if g.name in self.securityGroupNames]

    @property
    def _allSecurityGroups(self):
        return self.connection.get_all_security_groups()

    @property
    def connection(self):
        if not self._connection:
            self._connection = boto.ec2.connect_to_region(self.region)
        return self._connection

class CurrentHive(JsonConfigger):
    CONFIG = 'current.json'
    """configuration of an active bee hive (autogenerated)"""

    def __init__(self):
        super(CurrentHive, self).__init__(self.CONFIG)
        self.username = ''
        self.zone = ''
        self.beesIds = ''

    @property
    def isActive(self):
        return self._configPath.exists()

class Beekeeper(object):
    def __init__(self):
        self.bees = None

    @property
    def hive(self):
        if not hasattr(self, '_hive'):
            self._hive = CurrentHive()
        return self._hive

    @property
    def connectionWrapper(self):
        if not hasattr(self, '_connectionWrapper'):
            projectConfig = ProjectConfig()
            self._connectionWrapper = Ec2Connection(
                **projectConfig.get_config())
        return self._connectionWrapper

    @property
    def connection(self):
        if not hasattr(self, '_connection'):
            self._connection = self.connectionWrapper.connection
        return self._connection

    def _reserve_bees(self):
        if self.hive.isActive:
            log.warning("hive is up already: %s", self.hive.get_config())
            return

        if not self.connectionWrapper.keyPath.exists():
            raise BeeSting("key %s not found" % (self.connectionWrapper.keyPath))

        log.info(
            'attempting to call up %s bees' %
            (self.connectionWrapper.numberOfBees))
        reservation = self.connection.run_instances(
            image_id=self.connectionWrapper.instanceId,
            min_count=self.connectionWrapper.numberOfBees,
            max_count=self.connectionWrapper.numberOfBees,
            key_name=self.connectionWrapper.keyName,
            security_groups=self.connectionWrapper.securityGroupIds,
            instance_type=self.connectionWrapper.instanceType,
            placement=(None if self.connectionWrapper.isGovZone
                       else self.connectionWrapper.zone),
            subnet_id='')
        self._reservation = reservation

    def up(self):
        self._reserve_bees()
        while not self.allBeesAreUp:
            log.info('waiting for bees to load their machine guns... '
                     '%s bees are ready', self.activeBeesIds)
        self.connection.create_tags(self.activeBeesIds, {"Name": "a bee!"})
        self._hive.save_config()
        log.info('The swarm has assembled %i bees',
                 len(self._reservation.instances))

    @property
    def allBeesAreUp(self):
        return len(self.activeBeesIds) == self.connectionWrapper.numberOfBees

    @property
    def activeBeesIds(self):
        instanceIds = []
        for instance in self._reservation.instances:
            instance.update()
            if instance.state == 'running':
                instanceIds.append(instance.id)
        return instanceIds

class LoggingConfig(object):
    NAME = 'bees'

    def __init__(self):
        self.workPath = LocalWorkdir()
        self.localLogPath = None
        """:type: LocalPath"""

    def init_logging(self, logLevel=logging.INFO, logToFile=True):
        log.setLevel(logLevel)
        self.localLogPath = self.workPath / (self.NAME + '.log')
        fmt = '%(asctime)s %(name)s %(levelname)s: %(message)s'
        logging.basicConfig(format=fmt)
        if logToFile:
            fh = logging.FileHandler(filename=str(self.localLogPath))
            fh.setFormatter(logging.Formatter(fmt))
            log.addHandler(fh)
        log.name = self.NAME if log.name == '__main__' else log.name
        log.debug("working in %s", self.workPath)

def main(workingPath=None):
    workPath = LocalWorkdir()
    if workingPath:
        workPath.chdir(workingPath)
    log.info('working in %s', workPath)
    logCnf = LoggingConfig()
    logCnf.init_logging(logLevel=logging.DEBUG)
    cnf = ProjectConfig()

    print cnf.numberOfBees
    print cnf._config

    ec2 = Ec2Connection(**cnf._config)
    # for g in ec2._allSecurityGroups:
    #     print obj_attr(g)
    print ec2.securityGroupIds
    exit()

    print ec2.connection
    print ec2.__dict__.items()
    print type(ec2.connection)

    exit()
    hive = CurrentHive()
    print hive.isActive
    if hive.isActive:
        hive.load_config()
    print hive.isActive
    # print hive._config
    # print hive.isConfigured
    print hive._config

if __name__ == '__main__':
    main('../tests/fake_project_dir')

    # time.sleep(0.01)
    # for k, v in cnf.__dict__.items():
    #     print k, v
    # print cnf.beesIds
<|endoftext|>


<commit_before>import sys
import os
from benchexec import systeminfo
from elevate import elevate

import docker
from benchexec import benchexec

def init(config, benchmark):
    """
        This functions will set up the docker network to execute the test.
        As a result, it needs root permission for the setup part.
    """

    print("In P4Execution")
    client = docker.from_env()
    node1 = client.containers.run("alpine", "echo hello world", detach=True)
    node2 = client.containers.run("alpine", "echo hello world", detach=True)

def get_system_info():
    return systeminfo.SystemInfo()

def main(argv=None):
    if argv is None:
        argv = sys.argv

    print("Hello World")

if __name__ == "__main__":
    main()<commit_msg>p4executation now creates a simple setup<commit_after>import sys
import os
import logging
from benchexec import systeminfo
from elevate import elevate

import docker
from benchexec import benchexec

def init(config, benchmark):
    """
        This functions will set up the docker network to execute the test.
        As a result, it needs root permission for the setup part.
    """
    print("In P4Execution")

    client = docker.from_env()

    #Node containers
    node1 = client.containers.run("basic_node",
        command="python3 /usr/local/src/ptf/ptf_nn/ptf_nn_agent.py --device-socket 0@tcp://172.21.0.3:10001 -i 0-1@eth0",
        detach=True,
        name="node1",
        network="net1")
    node2 = client.containers.run("basic_node",
        command="python3 /usr/local/src/ptf/ptf_nn/ptf_nn_agent.py --device-socket 0@tcp://172.21.0.4:10002 -i 1-1@eth0",
        detach=True,
        name="node2",
        network="net2"
    )

    #Switch containers
    dir_path = os.path.dirname(os.path.realpath(__file__))
    mount1 = docker.types.Mount("/app", "/home/sdn/DockerTesting/Network_Compose/switch/switchDocs", type="bind")

    switch1 = client.containers.run("switch_test",
        command="python3 /app/ptf_nn_test_bridge.py -ifrom eth0 -ito eth1",
        detach=True,
        network="net1",
        name="switch1",
        mounts = [mount1]
        )

    ptf_tester = client.containers.run("ptf_tester",
        command="tail -f /dev/null",
        detach=True,
        network="mgnt",
        name="ptfTester")

def get_system_info():

    return systeminfo.SystemInfo()

def main(argv=None):
    if argv is None:
        argv = sys.argv

    print("Hello World")

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


<commit_before># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2016, Adrian Sampson.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

from __future__ import division, absolute_import, print_function
import six

"""Gets genres for imported music based on Last.fm tags.

Uses a provided whitelist file to determine which tags are valid genres.
The included (default) genre list was originally produced by scraping Wikipedia
and has been edited to remove some questionable entries.
The scraper script used is available here:
https://gist.github.com/1241307
"""
import pylast
import codecs
import os
import yaml
import traceback

from beets import plugins
from beets import ui
from beets import config
from beets.util import normpath, plurality
from beets import library

LASTFM = pylast.LastFMNetwork(api_key=plugins.LASTFM_KEY)

PYLAST_EXCEPTIONS = (
    pylast.WSError,
    pylast.MalformedResponseError,
    pylast.NetworkError,
)

REPLACE = {
    u'\u2010': '-',
}

def deduplicate(seq):
    """Remove duplicates from sequence wile preserving order.
    """
    seen = set()
    return [x for x in seq if x not in seen and not seen.add(x)]

# Canonicalization tree processing.

def flatten_tree(elem, path, branches):
    """Flatten nested lists/dictionaries into lists of strings
    (branches).
    """
    if not path:
        path = []

    if isinstance(elem, dict):
        for (k, v) in elem.items():
            flatten_tree(v, path + [k], branches)
    elif isinstance(elem, list):
        for sub in elem:
            flatten_tree(sub, path, branches)
    else:
        branches.append(path + [six.text_type(elem)])

def find_parents(candidate, branches):
    """Find parents genre of a given genre, ordered from the closest to
    the further parent.
    """
    for branch in branches:
        try:
            idx = branch.index(candidate.lower())
            return list(reversed(branch[:idx + 1]))
        except ValueError:
            continue
    return [candidate]

# Main plugin logic.

WHITELIST = os.path.join(os.path.dirname(__file__), 'genres.txt')
C14N_TREE = os.path.join(os.path.dirname(__file__), 'genres-tree.yaml')

class LastGenrePlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(LastGenrePlugin, self).__init__()

        self.config.add({
            'whitelist': True,
            'min_weight': 10,
            'count': 1,
            'fallback': None,
            'canonical': False,
            'source': 'album',
            'force': True,
            'auto': True,
            'separator': u', ',
            'specificity': False,
        })

        self.setup()

    def setup(self):
        """Setup plugin from config options
        """
        if self.config['auto']:
            self.import_stages = [self.imported]

        self._genre_cache = {}

        # Read the whitelist file if enabled.
        self.whitelist = set()
        wl_filename = self.config['whitelist'].get()
        if wl_filename in (True, ''):  # Indicates the default whitelist.
            wl_filename = WHITELIST
        if wl_filename:
            wl_filename = normpath(wl_filename)
            with open(wl_filename, 'rb') as f:
                for line in f:
                    line = line.decode('utf8').strip().lower()
                    if line and not line.startswith(u'#'):
                        self.whitelist.add(line)

        # Read the genres tree for canonicalization if enabled.
        self.c14n_branches = []
        c14n_filename = self.config['canonical'].get()
        if c14n_filename in (True, ''):  # Default tree.
            c14n_filename = C14N_TREE
        if c14n_filename:
            c14n_filename = normpath(c14n_filename)
            with codecs.open(c14n_filename, 'r', encoding='utf-8') as f:
                genres_tree = yaml.load(f)
            flatten_tree(genres_tree, [], self.c14n_branches)

    @property
    def sources(self):
        """A tuple of allowed genre sources. May contain 'track',
        'album', or 'artist.'
        """
        source = self.config['source'].as_choice(('track', 'album', 'artist'))
        if source == 'track':
            return 'track', 'album', 'artist'
        elif source == 'album':
            return 'album', 'artist'
        elif source == 'artist':
            return 'artist',

    def _get_depth(self, tag):
        """Find the depth of a tag in the genres tree.
        """
        depth = None
        for key, value in enumerate(self.c14n_branches):
            if tag in value:
                depth = value.index(tag)
                break
        return depth

    def _sort_by_depth(self, tags):
        """Given a list of tags, sort the tags by their depths in the
        genre tree.
        """
        tags_by_depth = []
        for key, value in enumerate(tags):
            depth = self._get_depth(value)
            tags_by_depth.append({'tag': value, 'depth': depth})
        tags_by_depth = sorted(tags_by_depth, key=lambda k: k['depth'], reverse=True)
        return [i['tag'] for i in tags_by_depth]

    def _resolve_genres(self, tags):
        """Given a list of strings, return a genre by joining them into a
        single string and (optionally) canonicalizing each.
        """
        if not tags:
            return None

        count = self.config['count'].get(int)
        if self.c14n_branches:
            # Extend the list to consider tags parents in the c14n tree
            tags_all = []
            for tag in tags:
                # Add parents that are in the whitelist, or add the oldest
                # ancestor if no whitelist
                if self.whitelist:
                    parents = [x for x in find_parents(tag, self.c14n_branches)
                               if self._is_allowed(x)]
                else:
                    parents = [find_parents(tag, self.c14n_branches)[-1]]

                tags_all += parents
                # Don't break here if we need to sort by specificity, which happens with
                # the full tag list below
                if not self.config['specificity'] and len(tags_all) >= count:
                    break
            tags = tags_all

        tags = deduplicate(tags)

        # sort the tags by specificity
        if self.config['specificity']:
            tags = self._sort_by_depth(tags)

        # c14n only adds allowed genres but we may have had forbidden genres in
        # the original tags list
        tags = [x.title() for x in tags if self._is_allowed(x)]

        return self.config['separator'].as_str().join(
            tags[:self.config['count'].get(int)]
        )

    def fetch_genre(self, lastfm_obj):
        """Return the genre for a pylast entity or None if no suitable genre
        can be found. Ex. 'Electronic, House, Dance'
        """
        min_weight = self.config['min_weight'].get(int)
        return self._resolve_genres(self._tags_for(lastfm_obj, min_weight))

    def _is_allowed(self, genre):
        """Determine whether the genre is present in the whitelist,
        returning a boolean.
        """
        if genre is None:
            return False
        if not self.whitelist or genre in self.whitelist:
            return True
        return False

    # Cached entity lookups.

    def _last_lookup(self, entity, method, *args):
        """Get a genre based on the named entity using the callable `method`
        whose arguments are given in the sequence `args`. The genre lookup
        is cached based on the entity name and the arguments. Before the
        lookup, each argument is has some Unicode characters replaced with
        rough ASCII equivalents in order to return better results from the
        Last.fm database.
        """
        # Shortcut if we're missing metadata.
        if any(not s for s in args):
            return None

        key = u'{0}.{1}'.format(entity,
                                u'-'.join(six.text_type(a) for a in args))
        if key in self._genre_cache:
            return self._genre_cache[key]
        else:
            args_replaced = []
            for arg in args:
                for k, v in REPLACE.items():
                    arg = arg.replace(k, v)
                args_replaced.append(arg)

            genre = self.fetch_genre(method(*args_replaced))
            self._genre_cache[key] = genre
            return genre

    def fetch_album_genre(self, obj):
        """Return the album genre for this Item or Album.
        """
        return self._last_lookup(
            u'album', LASTFM.get_album, obj.albumartist, obj.album
        )

    def fetch_album_artist_genre(self, obj):
        """Return the album artist genre for this Item or Album.
        """
        return self._last_lookup(
            u'artist', LASTFM.get_artist, obj.albumartist
        )

    def fetch_artist_genre(self, item):
        """Returns the track artist genre for this Item.
        """
        return self._last_lookup(
            u'artist', LASTFM.get_artist, item.artist
        )

    def fetch_track_genre(self, obj):
        """Returns the track genre for this Item.
        """
        return self._last_lookup(
            u'track', LASTFM.get_track, obj.artist, obj.title
        )

    def _get_genre(self, obj):
        """Get the genre string for an Album or Item object based on
        self.sources. Return a `(genre, source)` pair. The
        prioritization order is:
            - track (for Items only)
            - album
            - artist
            - original
            - fallback
            - None
        """

        # Shortcut to existing genre if not forcing.
        if not self.config['force'] and self._is_allowed(obj.genre):
            return obj.genre, 'keep'

        # Track genre (for Items only).
        if isinstance(obj, library.Item):
            if 'track' in self.sources:
                result = self.fetch_track_genre(obj)
                if result:
                    return result, 'track'

        # Album genre.
        if 'album' in self.sources:
            result = self.fetch_album_genre(obj)
            if result:
                return result, 'album'

        # Artist (or album artist) genre.
        if 'artist' in self.sources:
            result = None
            if isinstance(obj, library.Item):
                result = self.fetch_artist_genre(obj)
            elif obj.albumartist != config['va_name'].as_str():
                result = self.fetch_album_artist_genre(obj)
            else:
                # For "Various Artists", pick the most popular track genre.
                item_genres = []
                for item in obj.items():
                    item_genre = None
                    if 'track' in self.sources:
                        item_genre = self.fetch_track_genre(item)
                    if not item_genre:
                        item_genre = self.fetch_artist_genre(item)
                    if item_genre:
                        item_genres.append(item_genre)
                if item_genres:
                    result, _ = plurality(item_genres)

            if result:
                return result, 'artist'

        # Filter the existing genre.
        if obj.genre:
            result = self._resolve_genres([obj.genre])
            if result:
                return result, 'original'

        # Fallback string.
        fallback = self.config['fallback'].get()
        if fallback:
            return fallback, 'fallback'

        return None, None

    def commands(self):
        lastgenre_cmd = ui.Subcommand('lastgenre', help=u'fetch genres')
        lastgenre_cmd.parser.add_option(
            u'-f', u'--force', dest='force',
            action='store_true', default=False,
            help=u're-download genre when already present'
        )
        lastgenre_cmd.parser.add_option(
            u'-s', u'--source', dest='source', type='string',
            help=u'genre source: artist, album, or track'
        )

        def lastgenre_func(lib, opts, args):
            write = ui.should_write()
            self.config.set_args(opts)

            for album in lib.albums(ui.decargs(args)):
                album.genre, src = self._get_genre(album)
                self._log.info(u'genre for album {0} ({1}): {0.genre}',
                               album, src)
                album.store()

                for item in album.items():
                    # If we're using track-level sources, also look up each
                    # track on the album.
                    if 'track' in self.sources:
                        item.genre, src = self._get_genre(item)
                        item.store()
                        self._log.info(u'genre for track {0} ({1}): {0.genre}',
                                       item, src)

                    if write:
                        item.try_write()

        lastgenre_cmd.func = lastgenre_func
        return [lastgenre_cmd]

    def imported(self, session, task):
        """Event hook called when an import task finishes."""
        if task.is_album:
            album = task.album
            album.genre, src = self._get_genre(album)
            self._log.debug(u'added last.fm album genre ({0}): {1}',
                            src, album.genre)
            album.store()

            if 'track' in self.sources:
                for item in album.items():
                    item.genre, src = self._get_genre(item)
                    self._log.debug(u'added last.fm item genre ({0}): {1}',
                                    src, item.genre)
                    item.store()

        else:
            item = task.item
            item.genre, src = self._get_genre(item)
            self._log.debug(u'added last.fm item genre ({0}): {1}',
                            src, item.genre)
            item.store()

    def _tags_for(self, obj, min_weight=None):
        """Core genre identification routine.

        Given a pylast entity (album or track), return a list of
        tag names for that entity. Return an empty list if the entity is
        not found or another error occurs.

        If `min_weight` is specified, tags are filtered by weight.
        """
        # Work around an inconsistency in pylast where
        # Album.get_top_tags() does not return TopItem instances.
        # https://github.com/pylast/pylast/issues/86
        if isinstance(obj, pylast.Album):
            obj = super(pylast.Album, obj)

        try:
            res = obj.get_top_tags()
        except PYLAST_EXCEPTIONS as exc:
            self._log.debug(u'last.fm error: {0}', exc)
            return []
        except Exception as exc:
            # Isolate bugs in pylast.
            self._log.debug(u'{}', traceback.format_exc())
            self._log.error(u'error in pylast library: {0}', exc)
            return []

        # Filter by weight (optionally).
        if min_weight:
            res = [el for el in res if (int(el.weight or 0)) >= min_weight]

        # Get strings from tags.
        res = [el.item.get_name().lower() for el in res]

        return res
<commit_msg>better syntax; better comments<commit_after># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2016, Adrian Sampson.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

from __future__ import division, absolute_import, print_function
import six

"""Gets genres for imported music based on Last.fm tags.

Uses a provided whitelist file to determine which tags are valid genres.
The included (default) genre list was originally produced by scraping Wikipedia
and has been edited to remove some questionable entries.
The scraper script used is available here:
https://gist.github.com/1241307
"""
import pylast
import codecs
import os
import yaml
import traceback

from beets import plugins
from beets import ui
from beets import config
from beets.util import normpath, plurality
from beets import library

LASTFM = pylast.LastFMNetwork(api_key=plugins.LASTFM_KEY)

PYLAST_EXCEPTIONS = (
    pylast.WSError,
    pylast.MalformedResponseError,
    pylast.NetworkError,
)

REPLACE = {
    u'\u2010': '-',
}

def deduplicate(seq):
    """Remove duplicates from sequence wile preserving order.
    """
    seen = set()
    return [x for x in seq if x not in seen and not seen.add(x)]

# Canonicalization tree processing.

def flatten_tree(elem, path, branches):
    """Flatten nested lists/dictionaries into lists of strings
    (branches).
    """
    if not path:
        path = []

    if isinstance(elem, dict):
        for (k, v) in elem.items():
            flatten_tree(v, path + [k], branches)
    elif isinstance(elem, list):
        for sub in elem:
            flatten_tree(sub, path, branches)
    else:
        branches.append(path + [six.text_type(elem)])

def find_parents(candidate, branches):
    """Find parents genre of a given genre, ordered from the closest to
    the further parent.
    """
    for branch in branches:
        try:
            idx = branch.index(candidate.lower())
            return list(reversed(branch[:idx + 1]))
        except ValueError:
            continue
    return [candidate]

# Main plugin logic.

WHITELIST = os.path.join(os.path.dirname(__file__), 'genres.txt')
C14N_TREE = os.path.join(os.path.dirname(__file__), 'genres-tree.yaml')

class LastGenrePlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(LastGenrePlugin, self).__init__()

        self.config.add({
            'whitelist': True,
            'min_weight': 10,
            'count': 1,
            'fallback': None,
            'canonical': False,
            'source': 'album',
            'force': True,
            'auto': True,
            'separator': u', ',
            'specificity': False,
        })

        self.setup()

    def setup(self):
        """Setup plugin from config options
        """
        if self.config['auto']:
            self.import_stages = [self.imported]

        self._genre_cache = {}

        # Read the whitelist file if enabled.
        self.whitelist = set()
        wl_filename = self.config['whitelist'].get()
        if wl_filename in (True, ''):  # Indicates the default whitelist.
            wl_filename = WHITELIST
        if wl_filename:
            wl_filename = normpath(wl_filename)
            with open(wl_filename, 'rb') as f:
                for line in f:
                    line = line.decode('utf8').strip().lower()
                    if line and not line.startswith(u'#'):
                        self.whitelist.add(line)

        # Read the genres tree for canonicalization if enabled.
        self.c14n_branches = []
        c14n_filename = self.config['canonical'].get()
        if c14n_filename in (True, ''):  # Default tree.
            c14n_filename = C14N_TREE
        if c14n_filename:
            c14n_filename = normpath(c14n_filename)
            with codecs.open(c14n_filename, 'r', encoding='utf-8') as f:
                genres_tree = yaml.load(f)
            flatten_tree(genres_tree, [], self.c14n_branches)

    @property
    def sources(self):
        """A tuple of allowed genre sources. May contain 'track',
        'album', or 'artist.'
        """
        source = self.config['source'].as_choice(('track', 'album', 'artist'))
        if source == 'track':
            return 'track', 'album', 'artist'
        elif source == 'album':
            return 'album', 'artist'
        elif source == 'artist':
            return 'artist',

    def _get_depth(self, tag):
        """Find the depth of a tag in the genres tree.
        """
        depth = None
        for key, value in enumerate(self.c14n_branches):
            if tag in value:
                depth = value.index(tag)
                break
        return depth

    def _sort_by_depth(self, tags):
        """Given a list of tags, sort the tags by their depths in the
        genre tree.
        """
        depth_tag_pairs = [(self._get_depth(t), t) for t in tags]
        depth_tag_pairs.sort(reverse=True)
        return [p[1] for p in depth_tag_pairs]

    def _resolve_genres(self, tags):
        """Given a list of strings, return a genre by joining them into a
        single string and (optionally) canonicalizing each.
        """
        if not tags:
            return None

        count = self.config['count'].get(int)
        if self.c14n_branches:
            # Extend the list to consider tags parents in the c14n tree
            tags_all = []
            for tag in tags:
                # Add parents that are in the whitelist, or add the oldest
                # ancestor if no whitelist
                if self.whitelist:
                    parents = [x for x in find_parents(tag, self.c14n_branches)
                               if self._is_allowed(x)]
                else:
                    parents = [find_parents(tag, self.c14n_branches)[-1]]

                tags_all += parents
                # Stop if we have enough tags already, unless we need to find
                # the most specific tag (instead of the most popular).
                if not self.config['specificity'] and len(tags_all) >= count:
                    break
            tags = tags_all

        tags = deduplicate(tags)

        # Sort the tags by specificity.
        if self.config['specificity']:
            tags = self._sort_by_depth(tags)

        # c14n only adds allowed genres but we may have had forbidden genres in
        # the original tags list
        tags = [x.title() for x in tags if self._is_allowed(x)]

        return self.config['separator'].as_str().join(
            tags[:self.config['count'].get(int)]
        )

    def fetch_genre(self, lastfm_obj):
        """Return the genre for a pylast entity or None if no suitable genre
        can be found. Ex. 'Electronic, House, Dance'
        """
        min_weight = self.config['min_weight'].get(int)
        return self._resolve_genres(self._tags_for(lastfm_obj, min_weight))

    def _is_allowed(self, genre):
        """Determine whether the genre is present in the whitelist,
        returning a boolean.
        """
        if genre is None:
            return False
        if not self.whitelist or genre in self.whitelist:
            return True
        return False

    # Cached entity lookups.

    def _last_lookup(self, entity, method, *args):
        """Get a genre based on the named entity using the callable `method`
        whose arguments are given in the sequence `args`. The genre lookup
        is cached based on the entity name and the arguments. Before the
        lookup, each argument is has some Unicode characters replaced with
        rough ASCII equivalents in order to return better results from the
        Last.fm database.
        """
        # Shortcut if we're missing metadata.
        if any(not s for s in args):
            return None

        key = u'{0}.{1}'.format(entity,
                                u'-'.join(six.text_type(a) for a in args))
        if key in self._genre_cache:
            return self._genre_cache[key]
        else:
            args_replaced = []
            for arg in args:
                for k, v in REPLACE.items():
                    arg = arg.replace(k, v)
                args_replaced.append(arg)

            genre = self.fetch_genre(method(*args_replaced))
            self._genre_cache[key] = genre
            return genre

    def fetch_album_genre(self, obj):
        """Return the album genre for this Item or Album.
        """
        return self._last_lookup(
            u'album', LASTFM.get_album, obj.albumartist, obj.album
        )

    def fetch_album_artist_genre(self, obj):
        """Return the album artist genre for this Item or Album.
        """
        return self._last_lookup(
            u'artist', LASTFM.get_artist, obj.albumartist
        )

    def fetch_artist_genre(self, item):
        """Returns the track artist genre for this Item.
        """
        return self._last_lookup(
            u'artist', LASTFM.get_artist, item.artist
        )

    def fetch_track_genre(self, obj):
        """Returns the track genre for this Item.
        """
        return self._last_lookup(
            u'track', LASTFM.get_track, obj.artist, obj.title
        )

    def _get_genre(self, obj):
        """Get the genre string for an Album or Item object based on
        self.sources. Return a `(genre, source)` pair. The
        prioritization order is:
            - track (for Items only)
            - album
            - artist
            - original
            - fallback
            - None
        """

        # Shortcut to existing genre if not forcing.
        if not self.config['force'] and self._is_allowed(obj.genre):
            return obj.genre, 'keep'

        # Track genre (for Items only).
        if isinstance(obj, library.Item):
            if 'track' in self.sources:
                result = self.fetch_track_genre(obj)
                if result:
                    return result, 'track'

        # Album genre.
        if 'album' in self.sources:
            result = self.fetch_album_genre(obj)
            if result:
                return result, 'album'

        # Artist (or album artist) genre.
        if 'artist' in self.sources:
            result = None
            if isinstance(obj, library.Item):
                result = self.fetch_artist_genre(obj)
            elif obj.albumartist != config['va_name'].as_str():
                result = self.fetch_album_artist_genre(obj)
            else:
                # For "Various Artists", pick the most popular track genre.
                item_genres = []
                for item in obj.items():
                    item_genre = None
                    if 'track' in self.sources:
                        item_genre = self.fetch_track_genre(item)
                    if not item_genre:
                        item_genre = self.fetch_artist_genre(item)
                    if item_genre:
                        item_genres.append(item_genre)
                if item_genres:
                    result, _ = plurality(item_genres)

            if result:
                return result, 'artist'

        # Filter the existing genre.
        if obj.genre:
            result = self._resolve_genres([obj.genre])
            if result:
                return result, 'original'

        # Fallback string.
        fallback = self.config['fallback'].get()
        if fallback:
            return fallback, 'fallback'

        return None, None

    def commands(self):
        lastgenre_cmd = ui.Subcommand('lastgenre', help=u'fetch genres')
        lastgenre_cmd.parser.add_option(
            u'-f', u'--force', dest='force',
            action='store_true', default=False,
            help=u're-download genre when already present'
        )
        lastgenre_cmd.parser.add_option(
            u'-s', u'--source', dest='source', type='string',
            help=u'genre source: artist, album, or track'
        )

        def lastgenre_func(lib, opts, args):
            write = ui.should_write()
            self.config.set_args(opts)

            for album in lib.albums(ui.decargs(args)):
                album.genre, src = self._get_genre(album)
                self._log.info(u'genre for album {0} ({1}): {0.genre}',
                               album, src)
                album.store()

                for item in album.items():
                    # If we're using track-level sources, also look up each
                    # track on the album.
                    if 'track' in self.sources:
                        item.genre, src = self._get_genre(item)
                        item.store()
                        self._log.info(u'genre for track {0} ({1}): {0.genre}',
                                       item, src)

                    if write:
                        item.try_write()

        lastgenre_cmd.func = lastgenre_func
        return [lastgenre_cmd]

    def imported(self, session, task):
        """Event hook called when an import task finishes."""
        if task.is_album:
            album = task.album
            album.genre, src = self._get_genre(album)
            self._log.debug(u'added last.fm album genre ({0}): {1}',
                            src, album.genre)
            album.store()

            if 'track' in self.sources:
                for item in album.items():
                    item.genre, src = self._get_genre(item)
                    self._log.debug(u'added last.fm item genre ({0}): {1}',
                                    src, item.genre)
                    item.store()

        else:
            item = task.item
            item.genre, src = self._get_genre(item)
            self._log.debug(u'added last.fm item genre ({0}): {1}',
                            src, item.genre)
            item.store()

    def _tags_for(self, obj, min_weight=None):
        """Core genre identification routine.

        Given a pylast entity (album or track), return a list of
        tag names for that entity. Return an empty list if the entity is
        not found or another error occurs.

        If `min_weight` is specified, tags are filtered by weight.
        """
        # Work around an inconsistency in pylast where
        # Album.get_top_tags() does not return TopItem instances.
        # https://github.com/pylast/pylast/issues/86
        if isinstance(obj, pylast.Album):
            obj = super(pylast.Album, obj)

        try:
            res = obj.get_top_tags()
        except PYLAST_EXCEPTIONS as exc:
            self._log.debug(u'last.fm error: {0}', exc)
            return []
        except Exception as exc:
            # Isolate bugs in pylast.
            self._log.debug(u'{}', traceback.format_exc())
            self._log.error(u'error in pylast library: {0}', exc)
            return []

        # Filter by weight (optionally).
        if min_weight:
            res = [el for el in res if (int(el.weight or 0)) >= min_weight]

        # Get strings from tags.
        res = [el.item.get_name().lower() for el in res]

        return res
<|endoftext|>