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

import os
from oauth2client.file import Storage
from oauth2client.client import flow_from_clientsecrets
from apiclient.discovery import build

import httplib2
import pprint
import re

import logging
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

BASE_INFO = {
        "app":"gdcmdtools",
        "description":'Google Drive command line tools',
        "version":'1.03'}

GDAPI_VER = 'v2'
FTAPI_VER = 'v1'
DISCOVERY_URL = "https://www.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest"

DEBUG_LEVEL = ('debug', 'info', 'warning', 'error', 'critical')

class GDBase(object):
    def __init__(self):
        self.drive_service = None
        self.ft_service = None
        self.http = None
        self.root_folder = None

    def get_root(self):
        if self.root_folder == None:
            if self.drive_service == None:
                self.get_drive_service()
            about = self.drive_service.about().get().execute()

        self.root_folder = about['rootFolderId']
        logger.debug("root_folder=%s" % self.root_folder)
        return self.root_folder

    def get_drive_service(self, http):
        self.drive_service = build('drive', GDAPI_VER,
            discoveryServiceUrl=DISCOVERY_URL, http=http)
        return self.drive_service

    def get_ft_service(self, http):
        self.ft_service = build('fusiontables', FTAPI_VER,
            discoveryServiceUrl=DISCOVERY_URL, http=http)
        return self.ft_service

    @staticmethod
    def get_id_from_url(url):
        # normal, https://script.google.com/d/XXXXXXXXX/edit?usp=sharing
        normal_re = "^.*/d/([\w\-]*)"
        # folder, https://drive.google.com/folderview?id=XXXXXXXX&usp=sharing
        folder_share_re = "^.*/folderview\?id=([\w\-]*)"
        # folder on drive url: https://drive.google.com/drive/u/0/folders/XXXXXXXXXXX
        folder_url_re = "^.*/folders/([\w\-]*)"
        # open by id, https://drive.google.com/open?id=XXXXXXXXXX
        openbyid_re = "^.*/open\?id=([\w\-]*)"
        # webContentLink, https://docs.google.com/uc?id=XXXXXXXXXX&export=download
        webcontent_re = "^.*/uc\?id=([\w\-]*)"
        # export link, https://docs.google.com/feeds/download/documents/export/Export?id=XXXXXX&exportFormat=html
        export_re = "^.*/[Ee]xport\?id=([\w\-]*)"

        final_re = r"%s|%s|%s|%s|%s|%s" % \
                (normal_re, folder_share_re, folder_url_re, openbyid_re, webcontent_re, export_re)
        search = re.search(final_re, url)

        if(search):
            file_id = next(x for x in search.groups() if x is not None)
        else:
            file_id = url

        logger.debug("file_id: %s" % file_id)

        return file_id
<commit_msg>general re for id as parameter<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from oauth2client.file import Storage
from oauth2client.client import flow_from_clientsecrets
from apiclient.discovery import build

import httplib2
import pprint
import re

import logging
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

BASE_INFO = {
        "app":"gdcmdtools",
        "description":'Google Drive command line tools',
        "version":'1.03'}

GDAPI_VER = 'v2'
FTAPI_VER = 'v1'
DISCOVERY_URL = "https://www.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest"

DEBUG_LEVEL = ('debug', 'info', 'warning', 'error', 'critical')

class GDBase(object):
    def __init__(self):
        self.drive_service = None
        self.ft_service = None
        self.http = None
        self.root_folder = None

    def get_root(self):
        if self.root_folder == None:
            if self.drive_service == None:
                self.get_drive_service()
            about = self.drive_service.about().get().execute()

        self.root_folder = about['rootFolderId']
        logger.debug("root_folder=%s" % self.root_folder)
        return self.root_folder

    def get_drive_service(self, http):
        self.drive_service = build('drive', GDAPI_VER,
            discoveryServiceUrl=DISCOVERY_URL, http=http)
        return self.drive_service

    def get_ft_service(self, http):
        self.ft_service = build('fusiontables', FTAPI_VER,
            discoveryServiceUrl=DISCOVERY_URL, http=http)
        return self.ft_service

    @staticmethod
    def get_id_from_url(url):
        # normal, https://script.google.com/d/XXXXXXXXX/edit?usp=sharing
        normal_re = "^.*/d/([\w\-]*)"
        # folder on drive url: https://drive.google.com/drive/u/0/folders/XXXXXXXXXXX
        folder_url_re = "^.*/folders/([\w\-]*)"
        # folder, https://drive.google.com/folderview?id=XXXXXXXX&usp=sharing
        folder_share_re = "^.*/folderview\?id=([\w\-]*)"
        # open by id, https://drive.google.com/open?id=XXXXXXXXXX
        openbyid_re = "^.*/open\?id=([\w\-]*)"
        # webContentLink, https://docs.google.com/uc?id=XXXXXXXXXX&export=download
        webcontent_re = "^.*/uc\?id=([\w\-]*)"
        # export link, https://docs.google.com/feeds/download/documents/export/Export?id=XXXXXX&exportFormat=html
        export_re = "^.*/[Ee]xport\?id=([\w\-]*)"

        # general id as parameter
        id_parameter_re = "^.*\?id=([\w\-]*)"

        final_re = r"%s|%s|%s" % \
                (normal_re, folder_url_re, id_parameter_re)
                #(normal_re, folder_share_re, folder_url_re, openbyid_re, webcontent_re, export_re)
        search = re.search(final_re, url)

        if(search):
            file_id = next(x for x in search.groups() if x is not None)
        else:
            file_id = url

        logger.debug("file_id: %s" % file_id)

        return file_id
<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# pyglet
# Copyright (c) 2006-2008 Alex Holkner
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of pyglet nor the names of its
#    contributors may be used to endorse or promote products
#    derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ----------------------------------------------------------------------------

'''
'''

__docformat__ = 'restructuredtext'
__version__ = '$Id$'

from ctypes import *

from pyglet.gl import *
from pyglet.image import *
from pyglet.image.codecs import *
from pyglet.image.codecs import gif

import pyglet.lib
import pyglet.window

gdk = pyglet.lib.load_library('gdk-x11-2.0')
gdkpixbuf = pyglet.lib.load_library('gdk_pixbuf-2.0')

GdkPixbufLoader = c_void_p
GdkPixbuf = c_void_p
gdkpixbuf.gdk_pixbuf_loader_new.restype = GdkPixbufLoader
gdkpixbuf.gdk_pixbuf_loader_get_pixbuf.restype = GdkPixbuf
gdkpixbuf.gdk_pixbuf_get_pixels.restype = c_void_p
gdkpixbuf.gdk_pixbuf_loader_get_animation.restype = c_void_p
gdkpixbuf.gdk_pixbuf_animation_get_iter.restype = c_void_p
gdkpixbuf.gdk_pixbuf_animation_iter_get_pixbuf.restype = GdkPixbuf

class GTimeVal(Structure):
    _fields_ = [
        ('tv_sec', c_long),
        ('tv_usec', c_long)
    ]

class GdkPixbuf2ImageDecoder(ImageDecoder):
    def get_file_extensions(self):
        return ['.png', '.xpm', '.jpg', '.jpeg', '.tif', '.tiff', '.pnm',
                '.ras', '.bmp', '.gif']

    def get_animation_file_extensions(self):
        return ['.gif', '.ani']

    def _load(self, file, filename, load_func):
        data = file.read()
        loader = gdkpixbuf.gdk_pixbuf_loader_new()
        gdkpixbuf.gdk_pixbuf_loader_write(loader, data, len(data), None)
        if not gdkpixbuf.gdk_pixbuf_loader_close(loader, None):
            raise ImageDecodeException(filename)
        result = load_func(loader)
        if not result:
            raise ImageDecodeException('Unable to load: %s' % filename)
        return result

    def _pixbuf_to_image(self, pixbuf):
        # Get format and dimensions
        width = gdkpixbuf.gdk_pixbuf_get_width(pixbuf)
        height = gdkpixbuf.gdk_pixbuf_get_height(pixbuf)
        channels = gdkpixbuf.gdk_pixbuf_get_n_channels(pixbuf)
        rowstride = gdkpixbuf.gdk_pixbuf_get_rowstride(pixbuf)
        #has_alpha = gdkpixbuf.gdk_pixbuf_get_has_alpha(pixbuf)
        pixels = gdkpixbuf.gdk_pixbuf_get_pixels(pixbuf)

        # Copy pixel data.
        buffer = (c_ubyte * (rowstride * height))()
        memmove(buffer, pixels, rowstride * (height - 1) + width * channels)

        # Release pixbuf
        gdk.g_object_unref(pixbuf)

        # Determine appropriate GL type
        if channels == 3:
            format = 'RGB'
        else:
            format = 'RGBA'

        return ImageData(width, height, format, buffer, -rowstride)

    def decode(self, file, filename):
        pixbuf = self._load(file, filename,
                            gdkpixbuf.gdk_pixbuf_loader_get_pixbuf)

        return self._pixbuf_to_image(pixbuf)

    def decode_animation(self, file, filename):
        # Extract GIF control data.  If it's not a GIF, this method will
        # raise.
        gif_stream = gif.read(file)
        delays = [image.delay for image in gif_stream.images]

        # Get GDK animation iterator
        file.seek(0)
        anim = self._load(file, filename,
                          gdkpixbuf.gdk_pixbuf_loader_get_animation)
        time = GTimeVal(0, 0)
        iter = gdkpixbuf.gdk_pixbuf_animation_get_iter(anim, byref(time))

        frames = []

        # Extract each image
        for control_delay in delays:
            pixbuf = gdkpixbuf.gdk_pixbuf_animation_iter_get_pixbuf(iter)
            image = self._pixbuf_to_image(pixbuf)
            frames.append(AnimationFrame(image, control_delay))

            gdk_delay = gdkpixbuf.gdk_pixbuf_animation_iter_get_delay_time(iter)
            gdk_delay *= 1000 # milliseconds to microseconds
            # Compare gdk_delay to control_delay for interest only.
            #print control_delay, gdk_delay / 1000000.

            if gdk_delay == -1:
                break

            us = time.tv_usec + gdk_delay
            time.tv_sec += us // 1000000
            time.tv_usec = us % 1000000
            gdkpixbuf.gdk_pixbuf_animation_iter_advance(iter, byref(time))

        return Animation(frames)

def get_decoders():
    return [GdkPixbuf2ImageDecoder()]

def get_encoders():
    return []

def init():
    gdk.g_type_init()

init()

<commit_msg>Added Frankie.Robertson's fix for gif's with alpha. (Issue 411)<commit_after># ----------------------------------------------------------------------------
# pyglet
# Copyright (c) 2006-2008 Alex Holkner
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of pyglet nor the names of its
#    contributors may be used to endorse or promote products
#    derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ----------------------------------------------------------------------------

'''
'''

__docformat__ = 'restructuredtext'
__version__ = '$Id$'

from ctypes import *

from pyglet.gl import *
from pyglet.image import *
from pyglet.image.codecs import *
from pyglet.image.codecs import gif

import pyglet.lib
import pyglet.window

gdk = pyglet.lib.load_library('gdk-x11-2.0')
gdkpixbuf = pyglet.lib.load_library('gdk_pixbuf-2.0')

GdkPixbufLoader = c_void_p
GdkPixbuf = c_void_p
gdkpixbuf.gdk_pixbuf_loader_new.restype = GdkPixbufLoader
gdkpixbuf.gdk_pixbuf_loader_get_pixbuf.restype = GdkPixbuf
gdkpixbuf.gdk_pixbuf_get_pixels.restype = c_void_p
gdkpixbuf.gdk_pixbuf_loader_get_animation.restype = c_void_p
gdkpixbuf.gdk_pixbuf_animation_get_iter.restype = c_void_p
gdkpixbuf.gdk_pixbuf_animation_iter_get_pixbuf.restype = GdkPixbuf

class GTimeVal(Structure):
    _fields_ = [
        ('tv_sec', c_long),
        ('tv_usec', c_long)
    ]

class GdkPixbuf2ImageDecoder(ImageDecoder):
    def get_file_extensions(self):
        return ['.png', '.xpm', '.jpg', '.jpeg', '.tif', '.tiff', '.pnm',
                '.ras', '.bmp', '.gif']

    def get_animation_file_extensions(self):
        return ['.gif', '.ani']

    def _load(self, file, filename, load_func):
        data = file.read()
        loader = gdkpixbuf.gdk_pixbuf_loader_new()
        gdkpixbuf.gdk_pixbuf_loader_write(loader, data, len(data), None)
        if not gdkpixbuf.gdk_pixbuf_loader_close(loader, None):
            raise ImageDecodeException(filename)
        result = load_func(loader)
        if not result:
            raise ImageDecodeException('Unable to load: %s' % filename)
        return result

    def _pixbuf_to_image(self, pixbuf):
        # Get format and dimensions
        width = gdkpixbuf.gdk_pixbuf_get_width(pixbuf)
        height = gdkpixbuf.gdk_pixbuf_get_height(pixbuf)
        channels = gdkpixbuf.gdk_pixbuf_get_n_channels(pixbuf)
        rowstride = gdkpixbuf.gdk_pixbuf_get_rowstride(pixbuf)
        #has_alpha = gdkpixbuf.gdk_pixbuf_get_has_alpha(pixbuf)
        pixels = gdkpixbuf.gdk_pixbuf_get_pixels(pixbuf)

        # Copy pixel data.
        buffer = (c_ubyte * (rowstride * height))()
        memmove(buffer, pixels, rowstride * (height - 1) + width * channels)

        # Release pixbuf
        gdk.g_object_unref(pixbuf)

        # Determine appropriate GL type
        if channels == 3:
            format = 'RGB'
        else:
            format = 'RGBA'

        return ImageData(width, height, format, buffer, -rowstride)

    def decode(self, file, filename):
        pixbuf = self._load(file, filename,
                            gdkpixbuf.gdk_pixbuf_loader_get_pixbuf)

        return self._pixbuf_to_image(pixbuf)

    def decode_animation(self, file, filename):
        # Extract GIF control data.  If it's not a GIF, this method will
        # raise.
        gif_stream = gif.read(file)
        delays = [image.delay for image in gif_stream.images]

        # Get GDK animation iterator
        file.seek(0)
        anim = self._load(file, filename,
                          gdkpixbuf.gdk_pixbuf_loader_get_animation)
        time = GTimeVal(0, 0)
        iter = gdkpixbuf.gdk_pixbuf_animation_get_iter(anim, byref(time))

        frames = []

        # Extract each image
        for control_delay in delays:
            pixbuf = gdkpixbuf.gdk_pixbuf_animation_iter_get_pixbuf(iter)
            # When attempting to load animated gifs with an alpha channel on
            #  linux gdkpixbuf will normally return a null pixbuf for the final
            #  frame resulting in a segfault:
            #  http://code.google.com/p/pyglet/issues/detail?id=411
            # Since it is unclear why exactly this happens, the workaround
            #  below is to start again and extract that frame on its own.
            if pixbuf == None:
                file.seek(0)
                anim = self._load(file, filename,
                                gdkpixbuf.gdk_pixbuf_loader_get_animation)
                temptime = GTimeVal(0, 0)
                iter = gdkpixbuf.gdk_pixbuf_animation_get_iter(anim, byref(temptime))
                gdkpixbuf.gdk_pixbuf_animation_iter_advance(iter, byref(time))
                pixbuf = gdkpixbuf.gdk_pixbuf_animation_iter_get_pixbuf(iter)
            image = self._pixbuf_to_image(pixbuf)
            frames.append(AnimationFrame(image, control_delay))

            gdk_delay = gdkpixbuf.gdk_pixbuf_animation_iter_get_delay_time(iter)

            if gdk_delay == -1:
                break

            gdk_delay = gdkpixbuf.gdk_pixbuf_animation_iter_get_delay_time(iter)
            gdk_delay *= 1000 # milliseconds to microseconds
            # Compare gdk_delay to control_delay for interest only.
            #print control_delay, gdk_delay / 1000000.

            us = time.tv_usec + gdk_delay
            time.tv_sec += us // 1000000
            time.tv_usec = us % 1000000
            gdkpixbuf.gdk_pixbuf_animation_iter_advance(iter, byref(time))

        return Animation(frames)

def get_decoders():
    return [GdkPixbuf2ImageDecoder()]

def get_encoders():
    return []

def init():
    gdk.g_type_init()

init()

<|endoftext|>


<commit_before>#!/usr/bin/python
from __future__ import division

import numpy as np

from pymatgen.util.testing import PymatgenTest
from pymatgen.core.units import (Energy, Time, Length, unitized, Mass,
                                 EnergyArray, TimeArray, LengthArray, Unit,
                                 FloatWithUnit, ArrayWithUnit, UnitError)

import collections

class UnitTest(PymatgenTest):

    def test_init(self):
        u1 = Unit((("m", 1), ("s", -1)))
        self.assertEqual(str(u1), "m s^-1")
        u2 = Unit("kg m ^ 2 s ^ -2")
        self.assertEqual(str(u2), "J")
        self.assertEqual(str(u1 * u2), "J m s^-1")
        self.assertEqual(str(u2 / u1), "J s m^-1")
        self.assertEqual(str(u1 / Unit("m")), "s^-1")
        self.assertEqual(str(u1 * Unit("s")), "m")

        acc = u1 / Unit("s")
        newton = Unit("kg") * acc
        self.assertEqual(str(newton * Unit("m")), "N m")

class FloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        a = Energy(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = Energy(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        self.assertRaises(UnitError, Energy, 1, "m")

        d = Energy(1, "Ha")
        self.assertAlmostEqual(a + d, 28.31138386)
        self.assertAlmostEqual(a - d, -26.11138386)
        self.assertEqual(a + 1, 2.1)
        self.assertEqual(str(a / d), "1.1 eV Ha^-1")

    def test_time(self):
        a = Time(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60.0 h")
        self.assertEqual(str(3 * a), "60.0 h")

    def test_length(self):
        x = Length(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")
        self.assertEqual(str(x ** 3), "74.08800000000001 ang^3")

    def test_unitized(self):

        @unitized("eV")
        def f():
            return [1, 2, 3]

        self.assertEqual(str(f()[0]), "1.0 eV")
        self.assertIsInstance(f(), list)

        @unitized("eV")
        def g():
            return 2, 3, 4

        self.assertEqual(str(g()[0]), "2.0 eV")
        self.assertIsInstance(g(), tuple)

        @unitized("pm")
        def h():
            d = collections.OrderedDict()
            for i in range(3):
                d[i] = i * 20
            return d

        self.assertEqual(str(h()[1]), "20.0 pm")
        self.assertIsInstance(h(), collections.OrderedDict)

    def test_compound_operations(self):
        g = 10 * Length(1, "m") / (Time(1, "s") ** 2)
        e = Mass(1, "kg") * g * Length(1, "m")
        self.assertEqual(str(e), "10.0 N m")
        form_e = FloatWithUnit(10, unit="kJ mol^-1")
        self.assertEqual(str(form_e.to("eV atom^-1")), "0.10364269190469588 eV atom^-1")
        self.assertRaises(UnitError, form_e.to, "m s^-1")
        a = FloatWithUnit(1.0, "Ha^3")
        self.assertEqual(str(a.to("J^3")), "8.286726616151198e-53 J^3")
        a = FloatWithUnit(1.0, "Ha bohr^-2")
        self.assertEqual(str(a.to("J m^-2")), "1556.8929145749312 J m^-2")

class ArrayWithFloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        """
        Similar to FloatWithUnitTest.test_energy.
        Check whether EnergyArray and FloatWithUnit have same behavior.

        # TODO
        One can merge the two tests easily:

        for obj in [Energy, EnergyArray]:
            a = obj(...)
            self.assert(...)

        """
        a = EnergyArray(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = EnergyArray(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        # self.assertRaises(ValueError, Energy, 1, "m")

        d = EnergyArray(1, "Ha")
        self.assertAlmostEqual(a + d, 28.31138386)
        self.assertAlmostEqual(a - d, -26.11138386)
        self.assertEqual(a + 1, 2.1)

    def test_time(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        # here there's a minor difference because we have a ndarray with dtype=np.int.
        a = TimeArray(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60 h")
        self.assertEqual(str(3 * a), "60 h")

    def test_length(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        x = LengthArray(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")

    def test_array_algebra(self):
        ene_ha = EnergyArray([1, 2], "Ha")
        ene_ev = EnergyArray([1, 2], "eV")
        time_s = TimeArray([1, 2], "s")

        e1 = ene_ha.copy()
        e1 += 1
        e2 = ene_ha.copy()
        e2 -= 1
        e3 = ene_ha.copy()
        e3 /= 2
        e4 = ene_ha.copy()
        e4 *= 2

        objects_with_unit = [
            ene_ha + ene_ev,
            ene_ha - ene_ev,
            3 * ene_ha,
            ene_ha * 3,
            ene_ha / 3,
            3 / ene_ha,
            ene_ha * time_s,
            ene_ha / ene_ev,
            ene_ha.copy(),
            ene_ha[0:1],
            e1,
            e2,
            e3,
            e4,
        ]

        for i, obj in enumerate(objects_with_unit):
            #print(i, obj.unit)
            self.assertTrue(hasattr(obj, "unit"))
            #self.assertTrue(str(obj.unit) == "Ha")

        objects_without_unit = [
            # Here we could return a FloatWithUnit object but I prefer this
            # a bare scalar since FloatWithUnit extends float while we could have an int.
            ene_ha[0],
        ]

        for obj in objects_without_unit:
            self.assertFalse(hasattr(obj, "unit"))

        with self.assertRaises(UnitError):
            ene_ha + time_s

    def test_factors(self):
        e = EnergyArray([27.21138386, 1], "eV").to("Ha")
        self.assertTrue(str(e) == "[ 1.          0.03674933] Ha")
        l = LengthArray([1.0], "ang").to("bohr")
        self.assertTrue(str(l) == "[ 1.88972613] bohr")
        v = ArrayWithUnit([1,2,3], "bohr^3").to("ang^3")
        self.assertTrue(str(v) == '[ 0.14818471  0.29636942  0.44455413] '
                                  'ang^3')

if __name__ == '__main__':
    import unittest
    unittest.main()
<commit_msg>Fix unittests for Units.<commit_after>#!/usr/bin/python
from __future__ import division

import numpy as np

from pymatgen.util.testing import PymatgenTest
from pymatgen.core.units import (Energy, Time, Length, unitized, Mass,
                                 EnergyArray, TimeArray, LengthArray, Unit,
                                 FloatWithUnit, ArrayWithUnit, UnitError)

import collections

class UnitTest(PymatgenTest):

    def test_init(self):
        u1 = Unit((("m", 1), ("s", -1)))
        self.assertEqual(str(u1), "m s^-1")
        u2 = Unit("kg m ^ 2 s ^ -2")
        self.assertEqual(str(u2), "J")
        self.assertEqual(str(u1 * u2), "J m s^-1")
        self.assertEqual(str(u2 / u1), "J s m^-1")
        self.assertEqual(str(u1 / Unit("m")), "s^-1")
        self.assertEqual(str(u1 * Unit("s")), "m")

        acc = u1 / Unit("s")
        newton = Unit("kg") * acc
        self.assertEqual(str(newton * Unit("m")), "N m")

class FloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        a = Energy(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = Energy(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        self.assertRaises(UnitError, Energy, 1, "m")

        d = Energy(1, "Ha")
        self.assertAlmostEqual(a + d, 28.31138386)
        self.assertAlmostEqual(a - d, -26.11138386)
        self.assertEqual(a + 1, 2.1)
        self.assertEqual(str(a / d), "1.1 eV Ha^-1")

    def test_time(self):
        a = Time(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60.0 h")
        self.assertEqual(str(3 * a), "60.0 h")

    def test_length(self):
        x = Length(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")
        self.assertEqual(str(x ** 3), "74.088 ang^3")

    def test_unitized(self):

        @unitized("eV")
        def f():
            return [1, 2, 3]

        self.assertEqual(str(f()[0]), "1.0 eV")
        self.assertIsInstance(f(), list)

        @unitized("eV")
        def g():
            return 2, 3, 4

        self.assertEqual(str(g()[0]), "2.0 eV")
        self.assertIsInstance(g(), tuple)

        @unitized("pm")
        def h():
            d = collections.OrderedDict()
            for i in range(3):
                d[i] = i * 20
            return d

        self.assertEqual(str(h()[1]), "20.0 pm")
        self.assertIsInstance(h(), collections.OrderedDict)

    def test_compound_operations(self):
        g = 10 * Length(1, "m") / (Time(1, "s") ** 2)
        e = Mass(1, "kg") * g * Length(1, "m")
        self.assertEqual(str(e), "10.0 N m")
        form_e = FloatWithUnit(10, unit="kJ mol^-1")
        self.assertEqual(str(form_e.to("eV atom^-1")), "0.103642691905 eV atom^-1")
        self.assertRaises(UnitError, form_e.to, "m s^-1")
        a = FloatWithUnit(1.0, "Ha^3")
        self.assertEqual(str(a.to("J^3")), "8.28672661615e-53 J^3")
        a = FloatWithUnit(1.0, "Ha bohr^-2")
        self.assertEqual(str(a.to("J m^-2")), "1556.89291457 J m^-2")

class ArrayWithFloatWithUnitTest(PymatgenTest):

    def test_energy(self):
        """
        Similar to FloatWithUnitTest.test_energy.
        Check whether EnergyArray and FloatWithUnit have same behavior.

        # TODO
        One can merge the two tests easily:

        for obj in [Energy, EnergyArray]:
            a = obj(...)
            self.assert(...)

        """
        a = EnergyArray(1.1, "eV")
        b = a.to("Ha")
        self.assertAlmostEqual(b, 0.0404242579378)
        c = EnergyArray(3.14, "J")
        self.assertAlmostEqual(c.to("eV"), 1.9598339337836966e+19)
        # self.assertRaises(ValueError, Energy, 1, "m")

        d = EnergyArray(1, "Ha")
        self.assertAlmostEqual(a + d, 28.31138386)
        self.assertAlmostEqual(a - d, -26.11138386)
        self.assertEqual(a + 1, 2.1)

    def test_time(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        # here there's a minor difference because we have a ndarray with dtype=np.int.
        a = TimeArray(20, "h")
        self.assertAlmostEqual(a.to("s"), 3600 * 20)
        #Test left and right multiplication.
        self.assertEqual(str(a * 3), "60 h")
        self.assertEqual(str(3 * a), "60 h")

    def test_length(self):
        """
        Similar to FloatWithUnitTest.test_time.
        Check whether EnergyArray and FloatWithUnit have same behavior.
        """
        x = LengthArray(4.2, "ang")
        self.assertAlmostEqual(x.to("cm"), 4.2e-08)
        self.assertEqual(x.to("pm"), 420)
        self.assertEqual(str(x / 2), "2.1 ang")

    def test_array_algebra(self):
        ene_ha = EnergyArray([1, 2], "Ha")
        ene_ev = EnergyArray([1, 2], "eV")
        time_s = TimeArray([1, 2], "s")

        e1 = ene_ha.copy()
        e1 += 1
        e2 = ene_ha.copy()
        e2 -= 1
        e3 = ene_ha.copy()
        e3 /= 2
        e4 = ene_ha.copy()
        e4 *= 2

        objects_with_unit = [
            ene_ha + ene_ev,
            ene_ha - ene_ev,
            3 * ene_ha,
            ene_ha * 3,
            ene_ha / 3,
            3 / ene_ha,
            ene_ha * time_s,
            ene_ha / ene_ev,
            ene_ha.copy(),
            ene_ha[0:1],
            e1,
            e2,
            e3,
            e4,
        ]

        for i, obj in enumerate(objects_with_unit):
            #print(i, obj.unit)
            self.assertTrue(hasattr(obj, "unit"))
            #self.assertTrue(str(obj.unit) == "Ha")

        objects_without_unit = [
            # Here we could return a FloatWithUnit object but I prefer this
            # a bare scalar since FloatWithUnit extends float while we could have an int.
            ene_ha[0],
        ]

        for obj in objects_without_unit:
            self.assertFalse(hasattr(obj, "unit"))

        with self.assertRaises(UnitError):
            ene_ha + time_s

    def test_factors(self):
        e = EnergyArray([27.21138386, 1], "eV").to("Ha")
        self.assertTrue(str(e) == "[ 1.          0.03674933] Ha")
        l = LengthArray([1.0], "ang").to("bohr")
        self.assertTrue(str(l) == "[ 1.88972613] bohr")
        v = ArrayWithUnit([1,2,3], "bohr^3").to("ang^3")
        self.assertTrue(str(v) == '[ 0.14818471  0.29636942  0.44455413] '
                                  'ang^3')

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


<commit_before>"""cell decorator"""
import copy
import functools
import hashlib
import inspect
from typing import Any, Callable, Dict, Optional, Tuple

import toolz
from pydantic import BaseModel, validate_arguments

from gdsfactory.component import Component
from gdsfactory.name import MAX_NAME_LENGTH, clean_name, get_name_short
from gdsfactory.serialization import clean_dict, clean_value_name

CACHE: Dict[str, Component] = {}
INFO_VERSION = 2

class CellReturnTypeError(ValueError):
    pass

def clear_cache() -> None:
    """Clears Component CACHE."""
    global CACHE
    CACHE = {}

def print_cache():
    for k in CACHE:
        print(k)

def get_source_code(func: Callable) -> str:
    if isinstance(func, functools.partial):
        source = inspect.getsource(func.func)
    elif isinstance(func, toolz.functoolz.Compose):
        source = inspect.getsource(func.first)
    elif callable(func):
        source = inspect.getsource(func)
    else:
        raise ValueError(f"{func!r} needs to be callable")
    return source

class Settings(BaseModel):
    name: str
    module: str
    function_name: str

    info: Dict[str, Any]  # derived properties (length, resistance)
    info_version: int = INFO_VERSION

    full: Dict[str, Any]
    changed: Dict[str, Any]
    default: Dict[str, Any]

    child: Optional[Dict[str, Any]] = None

def cell_without_validator(func):
    """Decorator for Component functions.

    Similar to cell decorator, this one does not validate_arguments using
    type annotations

    I recommend using @cell instead
    """

    @functools.wraps(func)
    def _cell(*args, **kwargs):
        autoname = kwargs.pop("autoname", True)
        name = kwargs.pop("name", None)
        cache = kwargs.pop("cache", True)
        info = kwargs.pop("info", {})
        prefix = kwargs.pop("prefix", func.__name__)
        max_name_length = kwargs.pop("max_name_length", MAX_NAME_LENGTH)

        sig = inspect.signature(func)
        args_as_kwargs = dict(zip(sig.parameters.keys(), args))
        args_as_kwargs.update(**copy.deepcopy(kwargs))

        default = {
            p.name: p.default
            for p in sig.parameters.values()
            if not p.default == inspect._empty
        }
        changed = args_as_kwargs
        full = copy.deepcopy(default)
        full.update(**args_as_kwargs)

        default2 = copy.deepcopy(default)
        changed2 = copy.deepcopy(changed)

        # list of default args as strings
        default_args_list = [
            f"{key}={clean_value_name(default2[key])}" for key in sorted(default.keys())
        ]
        # list of explicitly passed args as strings
        passed_args_list = [
            f"{key}={clean_value_name(changed2[key])}" for key in sorted(changed.keys())
        ]

        # get only the args which are explicitly passed and different from defaults
        changed_arg_set = set(passed_args_list).difference(default_args_list)
        changed_arg_list = sorted(changed_arg_set)

        # if any args were different from default, append a hash of those args.
        # else, keep only the base name
        if changed_arg_list:
            named_args_string = "_".join(changed_arg_list)
            named_args_hash = hashlib.md5(named_args_string.encode()).hexdigest()[:8]
            name_signature = clean_name(f"{prefix}_{named_args_hash}")
        else:
            name_signature = prefix

        # filter the changed dictionary to only keep entries which have truly changed
        changed_arg_names = [carg.split("=")[0] for carg in changed_arg_list]
        changed = {k: changed[k] for k in changed_arg_names}

        name = name or name_signature
        decorator = kwargs.pop("decorator", None)
        name = get_name_short(name, max_name_length=max_name_length)

        if (
            "args" not in sig.parameters
            and "kwargs" not in sig.parameters
            and "settings" not in sig.parameters
        ):
            for key in kwargs.keys():
                if key not in sig.parameters.keys():
                    raise TypeError(
                        f"{func.__name__!r}() got invalid argument {key!r}\n"
                        f"valid arguments are {list(sig.parameters.keys())}"
                    )

        if cache and name in CACHE:
            # print(f"CACHE LOAD {name} {func.__name__}({arguments})")
            return CACHE[name]
        else:
            # print(f"BUILD {name} {func.__name__}({arguments})")
            assert callable(
                func
            ), f"{func!r} got decorated with @cell! @cell decorator is only for functions"

            component = func(*args, **kwargs)
            metadata_child = (
                dict(component.child.settings) if hasattr(component, "child") else None
            )

            if not isinstance(component, Component):
                raise CellReturnTypeError(
                    f"function {func.__name__!r} return type = {type(component)}",
                    "make sure that functions with @cell decorator return a Component",
                )

            if metadata_child and component.get_child_name:
                component_name = f"{metadata_child['name']}_{name}"
                component_name = get_name_short(
                    component_name, max_name_length=max_name_length
                )
            else:
                component_name = name

            if autoname:
                component.name = component_name

            component.info.update(**info)
            component.settings = Settings(
                name=component_name,
                module=func.__module__,
                function_name=func.__name__,
                changed=clean_dict(changed),
                default=clean_dict(default),
                full=clean_dict(full),
                info=component.info,
                child=metadata_child,
            )

            if decorator:
                if not callable(decorator):
                    raise ValueError(
                        f"decorator = {type(decorator)} needs to be callable"
                    )
                component_new = decorator(component)
                component = component_new or component

            component.lock()
            CACHE[name] = component
            return component

    return _cell

def cell(func, *args, **kwargs):
    """Decorator for Component functions.
    Wraps cell_without_validator Validates type annotations with pydantic.

    Implements a cache so that if a component has already been build
    it will return the component from the cache directly.
    This avoids 2 exact cells that are not references of the same cell
    You can always over-ride this with `cache = False`.

    When decorate your functions with @cell you get:

    - CACHE: avoids creating duplicated cells.
    - name: gives Components a unique name based on parameters.
    - adds Component.info with default, changed and full component settings.

    Keyword Args:
        autoname (bool): if True renames component based on args and kwargs
        name (str): Optional (ignored when autoname=True)
        cache (bool): returns component from the cache if it already exists.
            if False creates a new component
            by default True avoids having duplicated cells with the same name
        info: updates component.info dict
        prefix: name_prefix, defaults to function name
        max_name_length: truncates name beyond some characters (32) with a hash
        decorator: function to run over the component

    .. plot::
      :include-source:

      import gdsfactory as gf

      @gf.cell
      def rectangle(size=(4,2), layer=0)->gf.Component:
          c = gf.Component()
          w, h = size
          points = [[w, h], [w, 0], [0, 0], [0, h]]
          c.add_polygon(points, layer=layer)
          return c

      c = rectangle(layer=(1,0))
      c.plot()
    """
    return cell_without_validator(validate_arguments(func), *args, **kwargs)

@cell
def wg(length: int = 3, layer: Tuple[int, int] = (1, 0)) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    width = 0.5
    w = width / 2
    c.add_polygon([(0, -w), (length, -w), (length, w), (0, w)], layer=layer)
    c.add_port(name="o1", midpoint=[0, 0], width=width, orientation=180, layer=layer)
    c.add_port(name="o2", midpoint=[length, 0], width=width, orientation=0, layer=layer)
    return c

@cell
def wg2(wg1=wg) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    w = wg1()
    w1 = c << w
    w1.rotate(90)
    c.copy_child_info(w)
    c.add_ports(w1.ports)
    return c

@cell
def wg3(wg1=wg2) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    w = wg1()
    w1 = c << w
    w1.rotate(90)
    c.copy_child_info(w)
    c.add_ports(w1.ports)
    return c

def test_set_name() -> None:
    c = wg(length=3, name="hi_there")
    assert c.name == "hi_there", c.name

@cell
def demo(length: int = 3, wg_width: float = 0.5) -> Component:
    """Demo Dummy cell"""
    c = Component()
    w = length
    h = wg_width
    points = [
        [-w / 2.0, -h / 2.0],
        [-w / 2.0, h / 2],
        [w / 2, h / 2],
        [w / 2, -h / 2.0],
    ]
    c.add_polygon(points)
    return c

def test_names() -> None:
    name_base = demo().name
    assert name_base.split("_")[0] == "demo", name_base

    demo2 = functools.partial(demo, length=3)
    c1 = demo2(length=3)
    c2 = demo(length=3)
    assert c1.name == c2.name

    c1 = demo(length=3, wg_width=0.5).name
    c2 = demo(wg_width=0.5, length=3).name
    assert c1 == c2, f"{c1} != {c2}"

    name_with_prefix = demo(prefix="hi").name
    assert name_with_prefix.split("_")[0] == "hi", name_with_prefix

    name_args = demo(3).name
    name_kwargs = demo(length=3).name
    assert name_args == name_kwargs, name_with_prefix

    c1name = wg(length=3).name
    c2name = wg(length=3.0).name
    assert c1name == c2name

    # c1name = wg(length=3).name
    # c2name = wg(length=3.0).name
    # c3name = wg().name
    # assert c1name == c2name == c3name

    c = wg(length=3.1)
    assert c.settings.changed["length"] == 3.1

@cell
def straight_with_pins(**kwargs):
    import gdsfactory as gf

    c = gf.Component()
    ref = c << gf.components.straight()
    c.add_ports(ref.ports)
    gf.add_pins(c)
    return c

def test_import_gds_settings():
    """Sometimes it fails for files imported from GDS"""
    import gdsfactory as gf

    gdspath = gf.CONFIG["gdsdir"] / "mzi2x2.gds"
    c = gf.import_gds(gdspath)
    assert gf.routing.add_fiber_single(c)

if __name__ == "__main__":
    # test_names()
    # c = wg(layer=(1, 0))
    # print(c.info.changed)

    # import gdsfactory as gf
    # c = gf.components.straight()
    # c = gf.components.straight()
    # print(c.name)

    # c = wg3()
    # print(c.name)

    # print(wg(length=3).name)
    # print(wg(length=3.0).name)
    # print(wg().name)

    import gdsfactory as gf

    gdspath = gf.CONFIG["gdsdir"] / "mzi2x2.gds"
    c = gf.import_gds(gdspath)
    c3 = gf.routing.add_fiber_single(c)
    c3.show()
<commit_msg>replace assert with ValueError<commit_after>"""cell decorator"""
import copy
import functools
import hashlib
import inspect
from typing import Any, Callable, Dict, Optional, Tuple

import toolz
from pydantic import BaseModel, validate_arguments

from gdsfactory.component import Component
from gdsfactory.name import MAX_NAME_LENGTH, clean_name, get_name_short
from gdsfactory.serialization import clean_dict, clean_value_name

CACHE: Dict[str, Component] = {}
INFO_VERSION = 2

class CellReturnTypeError(ValueError):
    pass

def clear_cache() -> None:
    """Clears Component CACHE."""
    global CACHE
    CACHE = {}

def print_cache():
    for k in CACHE:
        print(k)

def get_source_code(func: Callable) -> str:
    if isinstance(func, functools.partial):
        source = inspect.getsource(func.func)
    elif isinstance(func, toolz.functoolz.Compose):
        source = inspect.getsource(func.first)
    elif callable(func):
        source = inspect.getsource(func)
    else:
        raise ValueError(f"{func!r} needs to be callable")
    return source

class Settings(BaseModel):
    name: str
    module: str
    function_name: str

    info: Dict[str, Any]  # derived properties (length, resistance)
    info_version: int = INFO_VERSION

    full: Dict[str, Any]
    changed: Dict[str, Any]
    default: Dict[str, Any]

    child: Optional[Dict[str, Any]] = None

def cell_without_validator(func):
    """Decorator for Component functions.

    Similar to cell decorator, this one does not validate_arguments using
    type annotations

    I recommend using @cell instead
    """

    @functools.wraps(func)
    def _cell(*args, **kwargs):
        autoname = kwargs.pop("autoname", True)
        name = kwargs.pop("name", None)
        cache = kwargs.pop("cache", True)
        info = kwargs.pop("info", {})
        prefix = kwargs.pop("prefix", func.__name__)
        max_name_length = kwargs.pop("max_name_length", MAX_NAME_LENGTH)

        sig = inspect.signature(func)
        args_as_kwargs = dict(zip(sig.parameters.keys(), args))
        args_as_kwargs.update(**copy.deepcopy(kwargs))

        default = {
            p.name: p.default
            for p in sig.parameters.values()
            if not p.default == inspect._empty
        }
        changed = args_as_kwargs
        full = copy.deepcopy(default)
        full.update(**args_as_kwargs)

        default2 = copy.deepcopy(default)
        changed2 = copy.deepcopy(changed)

        # list of default args as strings
        default_args_list = [
            f"{key}={clean_value_name(default2[key])}" for key in sorted(default.keys())
        ]
        # list of explicitly passed args as strings
        passed_args_list = [
            f"{key}={clean_value_name(changed2[key])}" for key in sorted(changed.keys())
        ]

        # get only the args which are explicitly passed and different from defaults
        changed_arg_set = set(passed_args_list).difference(default_args_list)
        changed_arg_list = sorted(changed_arg_set)

        # if any args were different from default, append a hash of those args.
        # else, keep only the base name
        if changed_arg_list:
            named_args_string = "_".join(changed_arg_list)
            named_args_hash = hashlib.md5(named_args_string.encode()).hexdigest()[:8]
            name_signature = clean_name(f"{prefix}_{named_args_hash}")
        else:
            name_signature = prefix

        # filter the changed dictionary to only keep entries which have truly changed
        changed_arg_names = [carg.split("=")[0] for carg in changed_arg_list]
        changed = {k: changed[k] for k in changed_arg_names}

        name = name or name_signature
        decorator = kwargs.pop("decorator", None)
        name = get_name_short(name, max_name_length=max_name_length)

        if (
            "args" not in sig.parameters
            and "kwargs" not in sig.parameters
            and "settings" not in sig.parameters
        ):
            for key in kwargs.keys():
                if key not in sig.parameters.keys():
                    raise TypeError(
                        f"{func.__name__!r}() got invalid argument {key!r}\n"
                        f"valid arguments are {list(sig.parameters.keys())}"
                    )

        if cache and name in CACHE:
            # print(f"CACHE LOAD {name} {func.__name__}({arguments})")
            return CACHE[name]
        else:
            # print(f"BUILD {name} {func.__name__}({arguments})")

            if not callable(func):
                raise ValueError(
                    f"{func!r} is not callable! @cell decorator is only for functions"
                )

            component = func(*args, **kwargs)
            metadata_child = (
                dict(component.child.settings) if hasattr(component, "child") else None
            )

            if not isinstance(component, Component):
                raise CellReturnTypeError(
                    f"function {func.__name__!r} return type = {type(component)}",
                    "make sure that functions with @cell decorator return a Component",
                )

            if metadata_child and component.get_child_name:
                component_name = f"{metadata_child['name']}_{name}"
                component_name = get_name_short(
                    component_name, max_name_length=max_name_length
                )
            else:
                component_name = name

            if autoname:
                component.name = component_name

            component.info.update(**info)
            component.settings = Settings(
                name=component_name,
                module=func.__module__,
                function_name=func.__name__,
                changed=clean_dict(changed),
                default=clean_dict(default),
                full=clean_dict(full),
                info=component.info,
                child=metadata_child,
            )

            if decorator:
                if not callable(decorator):
                    raise ValueError(
                        f"decorator = {type(decorator)} needs to be callable"
                    )
                component_new = decorator(component)
                component = component_new or component

            component.lock()
            CACHE[name] = component
            return component

    return _cell

def cell(func, *args, **kwargs):
    """Decorator for Component functions.
    Wraps cell_without_validator Validates type annotations with pydantic.

    Implements a cache so that if a component has already been build
    it will return the component from the cache directly.
    This avoids 2 exact cells that are not references of the same cell
    You can always over-ride this with `cache = False`.

    When decorate your functions with @cell you get:

    - CACHE: avoids creating duplicated cells.
    - name: gives Components a unique name based on parameters.
    - adds Component.info with default, changed and full component settings.

    Keyword Args:
        autoname (bool): if True renames component based on args and kwargs
        name (str): Optional (ignored when autoname=True)
        cache (bool): returns component from the cache if it already exists.
            if False creates a new component
            by default True avoids having duplicated cells with the same name
        info: updates component.info dict
        prefix: name_prefix, defaults to function name
        max_name_length: truncates name beyond some characters (32) with a hash
        decorator: function to run over the component

    .. plot::
      :include-source:

      import gdsfactory as gf

      @gf.cell
      def rectangle(size=(4,2), layer=0)->gf.Component:
          c = gf.Component()
          w, h = size
          points = [[w, h], [w, 0], [0, 0], [0, h]]
          c.add_polygon(points, layer=layer)
          return c

      c = rectangle(layer=(1,0))
      c.plot()
    """
    return cell_without_validator(validate_arguments(func), *args, **kwargs)

@cell
def wg(length: int = 3, layer: Tuple[int, int] = (1, 0)) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    width = 0.5
    w = width / 2
    c.add_polygon([(0, -w), (length, -w), (length, w), (0, w)], layer=layer)
    c.add_port(name="o1", midpoint=[0, 0], width=width, orientation=180, layer=layer)
    c.add_port(name="o2", midpoint=[length, 0], width=width, orientation=0, layer=layer)
    return c

@cell
def wg2(wg1=wg) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    w = wg1()
    w1 = c << w
    w1.rotate(90)
    c.copy_child_info(w)
    c.add_ports(w1.ports)
    return c

@cell
def wg3(wg1=wg2) -> Component:
    """Dummy component for testing."""
    from gdsfactory.component import Component

    c = Component("straight")
    w = wg1()
    w1 = c << w
    w1.rotate(90)
    c.copy_child_info(w)
    c.add_ports(w1.ports)
    return c

def test_set_name() -> None:
    c = wg(length=3, name="hi_there")
    assert c.name == "hi_there", c.name

@cell
def demo(length: int = 3, wg_width: float = 0.5) -> Component:
    """Demo Dummy cell"""
    c = Component()
    w = length
    h = wg_width
    points = [
        [-w / 2.0, -h / 2.0],
        [-w / 2.0, h / 2],
        [w / 2, h / 2],
        [w / 2, -h / 2.0],
    ]
    c.add_polygon(points)
    return c

def test_names() -> None:
    name_base = demo().name
    assert name_base.split("_")[0] == "demo", name_base

    demo2 = functools.partial(demo, length=3)
    c1 = demo2(length=3)
    c2 = demo(length=3)
    assert c1.name == c2.name

    c1 = demo(length=3, wg_width=0.5).name
    c2 = demo(wg_width=0.5, length=3).name
    assert c1 == c2, f"{c1} != {c2}"

    name_with_prefix = demo(prefix="hi").name
    assert name_with_prefix.split("_")[0] == "hi", name_with_prefix

    name_args = demo(3).name
    name_kwargs = demo(length=3).name
    assert name_args == name_kwargs, name_with_prefix

    c1name = wg(length=3).name
    c2name = wg(length=3.0).name
    assert c1name == c2name

    # c1name = wg(length=3).name
    # c2name = wg(length=3.0).name
    # c3name = wg().name
    # assert c1name == c2name == c3name

    c = wg(length=3.1)
    assert c.settings.changed["length"] == 3.1

@cell
def straight_with_pins(**kwargs):
    import gdsfactory as gf

    c = gf.Component()
    ref = c << gf.components.straight()
    c.add_ports(ref.ports)
    gf.add_pins(c)
    return c

def test_import_gds_settings():
    """Sometimes it fails for files imported from GDS"""
    import gdsfactory as gf

    gdspath = gf.CONFIG["gdsdir"] / "mzi2x2.gds"
    c = gf.import_gds(gdspath)
    assert gf.routing.add_fiber_single(c)

if __name__ == "__main__":
    # test_names()
    # c = wg(layer=(1, 0))
    # print(c.info.changed)

    # import gdsfactory as gf
    # c = gf.components.straight()
    # c = gf.components.straight()
    # print(c.name)

    # c = wg3()
    # print(c.name)

    # print(wg(length=3).name)
    # print(wg(length=3.0).name)
    # print(wg().name)

    import gdsfactory as gf

    gdspath = gf.CONFIG["gdsdir"] / "mzi2x2.gds"
    c = gf.import_gds(gdspath)
    c3 = gf.routing.add_fiber_single(c)
    c3.show()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function

from argparse import (
    ArgumentParser,
    Namespace,
    )
from datetime import datetime
import errno
import glob
import os
import re
import shutil
import subprocess
import sys
import tarfile
from urlparse import (
    urlsplit,
    urlunsplit,
    )

from debian import deb822

from agent_archive import get_agents
from build_package import juju_series
from make_agent_json import StanzaWriter
from utils import temp_dir

class NoDebsFound(Exception):
    """Raised when no .deb files could be found."""

# These are the archives that are searched for matching releases.
UBUNTU_ARCH = "http://archive.ubuntu.com/ubuntu/pool/universe/j/juju-core/"
ARM_ARCH = "http://ports.ubuntu.com/pool/universe/j/juju-core/"
PUBLIC_ARCHIVES = [UBUNTU_ARCH, ARM_ARCH]

def move_debs(dest_debs):
    juju_core_dir = os.path.join(dest_debs, 'juju-2.0')
    debs = glob.glob(os.path.join(juju_core_dir, '*deb'))
    if len(debs) == 0:
        # The juju-2.0 package was not found, try the juju-core package.
        print('No debs in {}'.format(juju_core_dir))
        juju_core_dir = os.path.join(dest_debs, 'juju-core')
        debs = glob.glob(os.path.join(juju_core_dir, '*deb'))
    if len(debs) == 0:
        print('No debs in {}'.format(juju_core_dir))
        raise NoDebsFound('No deb files found.')
    for deb in debs:
        shutil.move(deb, dest_debs)
    shutil.rmtree(juju_core_dir)

def retrieve_packages(release, upatch, archives, dest_debs, s3_config):
    # Retrieve the packages that contain a jujud for this version.
    print("Retrieving juju-core packages from archives")
    print(datetime.now().replace(microsecond=0).isoformat())
    for archive in archives:
        scheme, netloc, path, query, fragment = urlsplit(archive)
        # Strip username / password
        netloc = netloc.rsplit('@')[-1]
        safe_archive = urlunsplit((scheme, netloc, path, query, fragment))
        print("checking {} for {}".format(safe_archive, release))
        try:
            subprocess.call([
                'lftp', '-c', 'mirror', '-i',
                "(juju-2.0|juju-core).*{}.*\.{}~juj.*\.deb$".format(
                    release, upatch),
                archive], cwd=dest_debs)
        except OSError as e:
            if e.errno == errno.ENOENT:
                print("%s not found in %s" % (release, safe_archive))
            else:
                raise e
        except CalledProcessError:
            raise
    move_debs(dest_debs)
    if os.path.exists(s3_config):
        print(
            'checking s3://juju-qa-data/agent-archive for'
            ' {}.'.format(release))
        args = Namespace(
            version=release, destination=dest_debs, config=s3_config,
            dry_run=False, verbose=False)
        try:
            get_agents(args)
        except subprocess.CalledProcessError as e:
            print()
            sys.stderr.write(e.output)
            print("FAILED: get_agents() failed.")
            sys.exit(1)

def parse_args():
    parser = ArgumentParser()
    parser.add_argument('release', help='The juju release to prepare')
    parser.add_argument('agent_stream', help='The juju agent-stream.')
    parser.add_argument('destination', help='The simplestreams destination')
    parser.add_argument('--upatch', help='Ubuntu patchlevel', default='1')
    return parser.parse_args()

def list_ppas(juju_home):
    config = os.path.join(juju_home, 'buildarchrc')
    if not os.path.exists(config):
        return None
    listing = subprocess.check_output(
        ['/bin/bash', '-c', 'source {}; set -u; echo '
         '"$BUILD_STABLE1_ARCH\n'
         '$BUILD_DEVEL1_ARCH\n'
         '$BUILD_SUPPORTED_ARCH"'.format(config)])
    return listing.splitlines()

def deb_to_agent(deb_path, dest_dir, agent_stream):
    control_str = subprocess.check_output(['dpkg-deb', '-I', deb_path,
                                           'control'])
    control = deb822.Deb822(control_str)
    control_version = control['Version']
    base_version = re.sub('-0ubuntu.*$', '', control_version)
    major_minor = '.'.join(base_version.split('-')[0].split('.')[0:2])
    series = juju_series.get_name_from_package_version(control_version)
    architecture = control['Architecture']
    with temp_dir() as work_dir:
        contents = os.path.join(work_dir, 'contents')
        os.mkdir(contents)
        subprocess.check_call(['dpkg-deb', '-x', deb_path, contents])
        jujud_path = os.path.join(
            contents, 'usr', 'lib', 'juju-{}'.format(major_minor), 'bin',
            'jujud')
        if not os.path.exists(jujud_path):
            jujud_path = os.path.join(
                contents, 'usr', 'lib', 'juju-{}'.format(base_version), 'bin',
                'jujud')
        basename = 'juju-{}-{}-{}.tgz'.format(base_version, series,
                                              architecture)
        agent_filename = os.path.join(work_dir, basename)
        with tarfile.open(agent_filename, 'w:gz') as tf:
            tf.add(jujud_path, 'jujud')
        writer = StanzaWriter.for_ubuntu(
            juju_series.get_version(series), series, architecture,
            base_version, agent_filename, agent_stream=agent_stream)
        writer.write_stanzas()
        shutil.move(writer.filename, dest_dir)
        final_agent_path = os.path.join(dest_dir, writer.path)
        move_create_parent(agent_filename, final_agent_path)

def debs_to_agents(dest_debs, agent_stream):
    for deb_path in glob.glob(os.path.join(dest_debs, '*.deb')):
        deb_to_agent(deb_path, dest_debs, agent_stream)

def move_create_parent(source, target):
    try:
        os.makedirs(os.path.dirname(target))
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise
    shutil.move(source, target)

def make_windows_agent(dest_debs, agent_stream, release):
    source = os.path.join(
        dest_debs, 'juju-{}-win2012-amd64.tgz'.format(release))
    target = os.path.join(
        dest_debs, 'juju-{}-windows-amd64.tgz'.format(release))
    shutil.copy2(source, target)
    writer = StanzaWriter.for_windows(
        release, target, agent_stream=agent_stream)
    writer.write_stanzas()
    agent_path = os.path.join(dest_debs, writer.path)
    move_create_parent(target, agent_path)
    os.rename(writer.filename, os.path.join(dest_debs, writer.filename))

def make_centos_agent(dest_debs, agent_stream, release):
    tarfile = os.path.join(
        dest_debs, 'juju-{}-centos7-amd64.tgz'.format(release))
    writer = StanzaWriter.for_centos(release, tarfile,
                                     agent_stream=agent_stream)
    writer.write_stanzas()
    agent_path = os.path.join(dest_debs, writer.path)
    shutil.copy2(tarfile, agent_path)
    os.rename(writer.filename, os.path.join(dest_debs, writer.filename))

def main():
    args = parse_args()
    dest_debs = os.path.abspath(os.path.join(args.destination, 'debs'))
    juju_dir = os.environ.get(
        'JUJU_HOME', os.path.join(os.environ.get('HOME'), '.juju'))
    s3_config = os.path.join(juju_dir, 'juju-qa.s3cfg')
    archives = list_ppas(juju_dir)
    if archives is None:
        print("Only public archives will be searched.")
        archives = PUBLIC_ARCHIVES
    else:
        print("Searching the build archives.")
    retrieve_packages(args.release, args.upatch, archives, dest_debs,
                      s3_config)
    debs_to_agents(dest_debs, args.agent_stream)
    make_windows_agent(dest_debs, args.agent_stream, args.release)
    make_centos_agent(dest_debs, args.agent_stream, args.release)

if __name__ == '__main__':
    main()
<commit_msg>Fix error caught by exception.<commit_after>#!/usr/bin/env python
from __future__ import print_function

from argparse import (
    ArgumentParser,
    Namespace,
    )
from datetime import datetime
import errno
import glob
import os
import re
import shutil
import subprocess
import sys
import tarfile
from urlparse import (
    urlsplit,
    urlunsplit,
    )

from debian import deb822

from agent_archive import get_agents
from build_package import juju_series
from make_agent_json import StanzaWriter
from utils import temp_dir

class NoDebsFound(Exception):
    """Raised when no .deb files could be found."""

# These are the archives that are searched for matching releases.
UBUNTU_ARCH = "http://archive.ubuntu.com/ubuntu/pool/universe/j/juju-core/"
ARM_ARCH = "http://ports.ubuntu.com/pool/universe/j/juju-core/"
PUBLIC_ARCHIVES = [UBUNTU_ARCH, ARM_ARCH]

def move_debs(dest_debs):
    juju_core_dir = os.path.join(dest_debs, 'juju-2.0')
    debs = glob.glob(os.path.join(juju_core_dir, '*deb'))
    if len(debs) == 0:
        # The juju-2.0 package was not found, try the juju-core package.
        print('No debs in {}'.format(juju_core_dir))
        juju_core_dir = os.path.join(dest_debs, 'juju-core')
        debs = glob.glob(os.path.join(juju_core_dir, '*deb'))
    if len(debs) == 0:
        print('No debs in {}'.format(juju_core_dir))
        raise NoDebsFound('No deb files found.')
    for deb in debs:
        shutil.move(deb, dest_debs)
    shutil.rmtree(juju_core_dir)

def retrieve_packages(release, upatch, archives, dest_debs, s3_config):
    # Retrieve the packages that contain a jujud for this version.
    print("Retrieving juju-core packages from archives")
    print(datetime.now().replace(microsecond=0).isoformat())
    for archive in archives:
        scheme, netloc, path, query, fragment = urlsplit(archive)
        # Strip username / password
        netloc = netloc.rsplit('@')[-1]
        safe_archive = urlunsplit((scheme, netloc, path, query, fragment))
        print("checking {} for {}".format(safe_archive, release))
        try:
            subprocess.call([
                'lftp', '-c', 'mirror', '-i',
                "(juju-2.0|juju-core).*{}.*\.{}~juj.*\.deb$".format(
                    release, upatch),
                archive], cwd=dest_debs)
        except OSError as e:
            if e.errno == errno.ENOENT:
                print("%s not found in %s" % (release, safe_archive))
            else:
                raise e
        except subprocess.CalledProcessError:
            raise
    move_debs(dest_debs)
    if os.path.exists(s3_config):
        print(
            'checking s3://juju-qa-data/agent-archive for'
            ' {}.'.format(release))
        args = Namespace(
            version=release, destination=dest_debs, config=s3_config,
            dry_run=False, verbose=False)
        try:
            get_agents(args)
        except subprocess.CalledProcessError as e:
            print()
            sys.stderr.write(e.output)
            print("FAILED: get_agents() failed.")
            sys.exit(1)

def parse_args():
    parser = ArgumentParser()
    parser.add_argument('release', help='The juju release to prepare')
    parser.add_argument('agent_stream', help='The juju agent-stream.')
    parser.add_argument('destination', help='The simplestreams destination')
    parser.add_argument('--upatch', help='Ubuntu patchlevel', default='1')
    return parser.parse_args()

def list_ppas(juju_home):
    config = os.path.join(juju_home, 'buildarchrc')
    if not os.path.exists(config):
        return None
    listing = subprocess.check_output(
        ['/bin/bash', '-c', 'source {}; set -u; echo '
         '"$BUILD_STABLE1_ARCH\n'
         '$BUILD_DEVEL1_ARCH\n'
         '$BUILD_SUPPORTED_ARCH"'.format(config)])
    return listing.splitlines()

def deb_to_agent(deb_path, dest_dir, agent_stream):
    control_str = subprocess.check_output(['dpkg-deb', '-I', deb_path,
                                           'control'])
    control = deb822.Deb822(control_str)
    control_version = control['Version']
    base_version = re.sub('-0ubuntu.*$', '', control_version)
    major_minor = '.'.join(base_version.split('-')[0].split('.')[0:2])
    series = juju_series.get_name_from_package_version(control_version)
    architecture = control['Architecture']
    with temp_dir() as work_dir:
        contents = os.path.join(work_dir, 'contents')
        os.mkdir(contents)
        subprocess.check_call(['dpkg-deb', '-x', deb_path, contents])
        jujud_path = os.path.join(
            contents, 'usr', 'lib', 'juju-{}'.format(major_minor), 'bin',
            'jujud')
        if not os.path.exists(jujud_path):
            jujud_path = os.path.join(
                contents, 'usr', 'lib', 'juju-{}'.format(base_version), 'bin',
                'jujud')
        basename = 'juju-{}-{}-{}.tgz'.format(base_version, series,
                                              architecture)
        agent_filename = os.path.join(work_dir, basename)
        with tarfile.open(agent_filename, 'w:gz') as tf:
            tf.add(jujud_path, 'jujud')
        writer = StanzaWriter.for_ubuntu(
            juju_series.get_version(series), series, architecture,
            base_version, agent_filename, agent_stream=agent_stream)
        writer.write_stanzas()
        shutil.move(writer.filename, dest_dir)
        final_agent_path = os.path.join(dest_dir, writer.path)
        move_create_parent(agent_filename, final_agent_path)

def debs_to_agents(dest_debs, agent_stream):
    for deb_path in glob.glob(os.path.join(dest_debs, '*.deb')):
        deb_to_agent(deb_path, dest_debs, agent_stream)

def move_create_parent(source, target):
    try:
        os.makedirs(os.path.dirname(target))
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise
    shutil.move(source, target)

def make_windows_agent(dest_debs, agent_stream, release):
    source = os.path.join(
        dest_debs, 'juju-{}-win2012-amd64.tgz'.format(release))
    target = os.path.join(
        dest_debs, 'juju-{}-windows-amd64.tgz'.format(release))
    shutil.copy2(source, target)
    writer = StanzaWriter.for_windows(
        release, target, agent_stream=agent_stream)
    writer.write_stanzas()
    agent_path = os.path.join(dest_debs, writer.path)
    move_create_parent(target, agent_path)
    os.rename(writer.filename, os.path.join(dest_debs, writer.filename))

def make_centos_agent(dest_debs, agent_stream, release):
    tarfile = os.path.join(
        dest_debs, 'juju-{}-centos7-amd64.tgz'.format(release))
    writer = StanzaWriter.for_centos(release, tarfile,
                                     agent_stream=agent_stream)
    writer.write_stanzas()
    agent_path = os.path.join(dest_debs, writer.path)
    shutil.copy2(tarfile, agent_path)
    os.rename(writer.filename, os.path.join(dest_debs, writer.filename))

def main():
    args = parse_args()
    dest_debs = os.path.abspath(os.path.join(args.destination, 'debs'))
    juju_dir = os.environ.get(
        'JUJU_HOME', os.path.join(os.environ.get('HOME'), '.juju'))
    s3_config = os.path.join(juju_dir, 'juju-qa.s3cfg')
    archives = list_ppas(juju_dir)
    if archives is None:
        print("Only public archives will be searched.")
        archives = PUBLIC_ARCHIVES
    else:
        print("Searching the build archives.")
    retrieve_packages(args.release, args.upatch, archives, dest_debs,
                      s3_config)
    debs_to_agents(dest_debs, args.agent_stream)
    make_windows_agent(dest_debs, args.agent_stream, args.release)
    make_centos_agent(dest_debs, args.agent_stream, args.release)

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


<commit_before>#!/usr/bin/env python

from time import sleep, strptime
import pytest
import yaml
from subprocess import check_output
from os.path import isdir, dirname, isfile, join
from shutil import rmtree
from epiphan import parse, dump, BGR24, I420
import pygiftgrab

recorders = []

def frame_rate(file_path):
    """Use ``ffprobe`` to inspect frame rate of video file

    @param file_path
    @return
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if frame rate cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=r_frame_rate',
                                   file_path])
    frame_rate_str = ffprobe_output.split('=')[1].strip()[1:-1].split('/')
    fr = -1
    if len(frame_rate_str) == 1:
        fr = float(frame_rate_str[0])
    elif len(frame_rate_str) == 2:
        fr = float(frame_rate_str[0])/float(frame_rate_str[1])
    else:
        raise ValueError('Frame rate cannot be inferred from ffprobe output')

    return fr

def duration(file_path):
    """Use ``ffprobe`` to inspect duration of video file

    @param file_path
    @return
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if duration cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=duration',
                                   file_path])
    duration_str = ffprobe_output.split('=')[1].strip()[1:-1]
    dur = -1
    if duration_str is not None:
        dur = float(duration_str)
    else:
        raise ValueError('Duration cannot be inferred from ffprobe output')

    return dur

def resolution(file_path):
    """Use ``ffprobe`` to inspect resolution of video file

    @param file_path
    @return width, height
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if resolution cannot be inferred from
    ``ffprobe`` output
    """
    width = -1
    height = -1
    ffprobe_commands = [ ['ffprobe',
                          '-v', 'error',
                          '-of', 'flat=s=_',
                          '-show_entries', 'stream=width',
                          file_path],
                         ['ffprobe',
                          '-v', 'error',
                          '-of', 'flat=s=_',
                          '-show_entries', 'stream=height',
                          file_path] ]
    res = []
    for ffprobe_command in ffprobe_commands:
        ffprobe_output = check_output(ffprobe_command)
        resolution_str = ffprobe_output.split('=')[1].strip()
        if resolution_str is not None:
            res.append(float(resolution_str))
        else:
            raise ValueError('Resolution cannot be inferred from ffprobe output')
    width = res[0]
    height = res[1]

    return width, height

def codec(file_path):
    """Use ``ffprobe`` to inspect codec of video file

    @param file_path
    @return width, height
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if codec cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=codec_name',
                                   file_path])
    codec_str = ffprobe_output.split('=')[1].strip().strip('"')
    if codec_str is not None:
        return codec_str
    else:
        raise ValueError('Codec cannot be inferred from ffprobe output')

def timing_report(file_path):
    """Check whether `file_path` is a valid YAML timing report file

    @param file_path
    @return ``True`` if valid, ``False`` otherwise
    @throw IOError if `file_path` cannot be opened for reading
    @throw YAMLError if `file_path` cannot be parsed as a YAML file
    """
    with open(file_path, 'r') as stream:
        data = yaml.load(stream)
        if data['elapsed'] is None:
            return False
        try:
            strptime(data['elapsed'], '%H:%M:%S.%f')
        except ValueError:
            return False
        if data['latency'] is None:
            return False
        try:
            strptime(data['latency'], '%H:%M:%S.%f')
        except ValueError:
            return False
    return True

def test_parse(colour_space):
    # not existing config file
    with pytest.raises(IOError):
        _ = parse('/this/file/should/never/exist.yml')

    # non-parseable file
    with pytest.raises(yaml.YAMLError):
        _ = parse('config/yamlerror.yml')

    # files with invalid values
    with pytest.raises(ValueError):
        _ = parse('config/valueerror1.yml')
    if colour_space == BGR24:
        with pytest.raises(ValueError):
            _ = parse('config/valueerror2.yml')
    elif colour_space == I420:
        with pytest.raises(ValueError):
            _ = parse('config/valueerror3.yml')

    # folder that can't be created
    with pytest.raises(OSError):
        _ = parse('config/oserror.yml')

def test_frame_grabbing(colour_space):
    # test-input & data
    if colour_space == BGR24:
        fs_config_file = 'config/sdi.yml'
        us_config_file = 'config/dvi.yml'
        fs_frame_rate = 28.0
        us_frame_rate = 14.0
    elif colour_space == I420:
        fs_config_file = 'config/sdi-i420.yml'
        us_config_file = 'config/dvi-i420.yml'
        fs_frame_rate = 56.0
        us_frame_rate = 28.0
    else:
        pytest.fail('Colour space not set')

    # create actual threads
    fs = parse(fs_config_file)
    us = parse(us_config_file)
    assert fs is not None
    recorders.append(fs)
    assert us is not None
    recorders.append(us)
    assert fs.frame_rate == fs_frame_rate and \
        fs.port == pygiftgrab.Device.DVI2PCIeDuo_SDI and \
        fs.timeout_limit == 20.0 and \
        fs.colour_space == colour_space
    assert us.frame_rate == us_frame_rate and \
        us.port == pygiftgrab.Device.DVI2PCIeDuo_DVI and \
        us.timeout_limit == 15.0 and \
        us.colour_space == colour_space

    # test actual output now
    roi = [426, 40, 1068, 1040]
    fs.start()
    us.start()
    fs.set_sub_frame(roi[0], roi[1], roi[2], roi[3])  # should have no effect
    recording_duration = 5
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)

    # has output directory been created?
    assert isdir(dirname(fs.file_path))
    assert isdir(dirname(us.file_path))

    # have videos been recorded?
    fs_file_path = fs.file_path + '-000001.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000001.mp4'
    assert timing_report(us_file_path + '.timing.yml')
    assert isfile(us_file_path)
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (1920, 1080)
    assert resolution(us_file_path) == (1920, 1080)
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # another recording
    with pytest.raises(ValueError):
        fs.set_sub_frame(roi[0], -roi[1], roi[2], roi[3])
    fs.set_sub_frame(roi[0], roi[1], roi[2], roi[3])
    us.set_sub_frame(roi[0], roi[1], roi[2], roi[3])
    recording_duration = 10
    fs.resume_recording()
    us.resume_recording()
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)
    fs_file_path = fs.file_path + '-000002.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000002.mp4'
    assert isfile(us_file_path)
    assert timing_report(us_file_path + '.timing.yml')
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (roi[2], roi[3])
    assert resolution(us_file_path) == (roi[2], roi[3])
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # yet another recording
    recording_duration = 25
    fs.set_full_frame()
    fs.resume_recording()
    us.resume_recording()
    us.set_full_frame()  # should have no effect
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)
    fs_file_path = fs.file_path + '-000003.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000003.mp4'
    assert timing_report(us_file_path + '.timing.yml')
    assert isfile(us_file_path)
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (1920, 1080)
    assert resolution(us_file_path) == (roi[2], roi[3])
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # finish session
    fs.stop()
    us.stop()

    # join threads
    fs.join(timeout=fs.timeout_limit)
    assert not fs.isAlive()
    us.join(timeout=us.timeout_limit)
    assert not us.isAlive()

    # test dumped Recorder configuration
    dump(fs)
    dumped_fs = parse(join(dirname(fs_file_path), 'config.yml'))
    assert dirname(dirname(dumped_fs.file_path)) == dirname(fs.file_path)
    assert dumped_fs.frame_rate == fs.frame_rate
    assert dumped_fs.port == fs.port
    assert dumped_fs.timeout_limit == fs.timeout_limit
    dump(us)
    dumped_us = parse(join(dirname(us_file_path), 'config.yml'))
    assert dirname(dirname(dumped_us.file_path)) == dirname(us.file_path)
    assert dumped_us.frame_rate == us.frame_rate
    assert dumped_us.port == us.port
    assert dumped_us.timeout_limit == us.timeout_limit

@pytest.yield_fixture(autouse=True)
def cleanup():
    # This section runs before each test

    # Run test
    yield
    # This section runs after each test
    for recorder in recorders:
        recorder.stop()
    sleep(5)

    for recorder in recorders:
        recorder.join(timeout=recorder.timeout_limit)
    sleep(5)

    for recorder in recorders:
        rmtree(dirname(recorder.file_path))

    for recorder in recorders:
        assert not recorder.isAlive()
<commit_msg>Issue #48: failing if colour space not set in test_parse<commit_after>#!/usr/bin/env python

from time import sleep, strptime
import pytest
import yaml
from subprocess import check_output
from os.path import isdir, dirname, isfile, join
from shutil import rmtree
from epiphan import parse, dump, BGR24, I420
import pygiftgrab

recorders = []

def frame_rate(file_path):
    """Use ``ffprobe`` to inspect frame rate of video file

    @param file_path
    @return
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if frame rate cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=r_frame_rate',
                                   file_path])
    frame_rate_str = ffprobe_output.split('=')[1].strip()[1:-1].split('/')
    fr = -1
    if len(frame_rate_str) == 1:
        fr = float(frame_rate_str[0])
    elif len(frame_rate_str) == 2:
        fr = float(frame_rate_str[0])/float(frame_rate_str[1])
    else:
        raise ValueError('Frame rate cannot be inferred from ffprobe output')

    return fr

def duration(file_path):
    """Use ``ffprobe`` to inspect duration of video file

    @param file_path
    @return
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if duration cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=duration',
                                   file_path])
    duration_str = ffprobe_output.split('=')[1].strip()[1:-1]
    dur = -1
    if duration_str is not None:
        dur = float(duration_str)
    else:
        raise ValueError('Duration cannot be inferred from ffprobe output')

    return dur

def resolution(file_path):
    """Use ``ffprobe`` to inspect resolution of video file

    @param file_path
    @return width, height
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if resolution cannot be inferred from
    ``ffprobe`` output
    """
    width = -1
    height = -1
    ffprobe_commands = [ ['ffprobe',
                          '-v', 'error',
                          '-of', 'flat=s=_',
                          '-show_entries', 'stream=width',
                          file_path],
                         ['ffprobe',
                          '-v', 'error',
                          '-of', 'flat=s=_',
                          '-show_entries', 'stream=height',
                          file_path] ]
    res = []
    for ffprobe_command in ffprobe_commands:
        ffprobe_output = check_output(ffprobe_command)
        resolution_str = ffprobe_output.split('=')[1].strip()
        if resolution_str is not None:
            res.append(float(resolution_str))
        else:
            raise ValueError('Resolution cannot be inferred from ffprobe output')
    width = res[0]
    height = res[1]

    return width, height

def codec(file_path):
    """Use ``ffprobe`` to inspect codec of video file

    @param file_path
    @return width, height
    @throw OSError if `file_path` does not exist, or if
    ``ffprobe`` cannot be run (e.g. not installed)
    @throw ValueError if codec cannot be inferred from
    ``ffprobe`` output
    """
    ffprobe_output = check_output(['ffprobe',
                                   '-v', 'error',
                                   '-of', 'flat=s=_',
                                   '-show_entries', 'stream=codec_name',
                                   file_path])
    codec_str = ffprobe_output.split('=')[1].strip().strip('"')
    if codec_str is not None:
        return codec_str
    else:
        raise ValueError('Codec cannot be inferred from ffprobe output')

def timing_report(file_path):
    """Check whether `file_path` is a valid YAML timing report file

    @param file_path
    @return ``True`` if valid, ``False`` otherwise
    @throw IOError if `file_path` cannot be opened for reading
    @throw YAMLError if `file_path` cannot be parsed as a YAML file
    """
    with open(file_path, 'r') as stream:
        data = yaml.load(stream)
        if data['elapsed'] is None:
            return False
        try:
            strptime(data['elapsed'], '%H:%M:%S.%f')
        except ValueError:
            return False
        if data['latency'] is None:
            return False
        try:
            strptime(data['latency'], '%H:%M:%S.%f')
        except ValueError:
            return False
    return True

def test_parse(colour_space):
    # not existing config file
    with pytest.raises(IOError):
        _ = parse('/this/file/should/never/exist.yml')

    # non-parseable file
    with pytest.raises(yaml.YAMLError):
        _ = parse('config/yamlerror.yml')

    # files with invalid values
    with pytest.raises(ValueError):
        _ = parse('config/valueerror1.yml')
    if colour_space == BGR24:
        with pytest.raises(ValueError):
            _ = parse('config/valueerror2.yml')
    elif colour_space == I420:
        with pytest.raises(ValueError):
            _ = parse('config/valueerror3.yml')
    else:
        pytest.fail('Colour space not configured properly')

    # folder that can't be created
    with pytest.raises(OSError):
        _ = parse('config/oserror.yml')

def test_frame_grabbing(colour_space):
    # test-input & data
    if colour_space == BGR24:
        fs_config_file = 'config/sdi.yml'
        us_config_file = 'config/dvi.yml'
        fs_frame_rate = 28.0
        us_frame_rate = 14.0
    elif colour_space == I420:
        fs_config_file = 'config/sdi-i420.yml'
        us_config_file = 'config/dvi-i420.yml'
        fs_frame_rate = 56.0
        us_frame_rate = 28.0
    else:
        pytest.fail('Colour space not configured properly')

    # create actual threads
    fs = parse(fs_config_file)
    us = parse(us_config_file)
    assert fs is not None
    recorders.append(fs)
    assert us is not None
    recorders.append(us)
    assert fs.frame_rate == fs_frame_rate and \
        fs.port == pygiftgrab.Device.DVI2PCIeDuo_SDI and \
        fs.timeout_limit == 20.0 and \
        fs.colour_space == colour_space
    assert us.frame_rate == us_frame_rate and \
        us.port == pygiftgrab.Device.DVI2PCIeDuo_DVI and \
        us.timeout_limit == 15.0 and \
        us.colour_space == colour_space

    # test actual output now
    roi = [426, 40, 1068, 1040]
    fs.start()
    us.start()
    fs.set_sub_frame(roi[0], roi[1], roi[2], roi[3])  # should have no effect
    recording_duration = 5
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)

    # has output directory been created?
    assert isdir(dirname(fs.file_path))
    assert isdir(dirname(us.file_path))

    # have videos been recorded?
    fs_file_path = fs.file_path + '-000001.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000001.mp4'
    assert timing_report(us_file_path + '.timing.yml')
    assert isfile(us_file_path)
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (1920, 1080)
    assert resolution(us_file_path) == (1920, 1080)
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # another recording
    with pytest.raises(ValueError):
        fs.set_sub_frame(roi[0], -roi[1], roi[2], roi[3])
    fs.set_sub_frame(roi[0], roi[1], roi[2], roi[3])
    us.set_sub_frame(roi[0], roi[1], roi[2], roi[3])
    recording_duration = 10
    fs.resume_recording()
    us.resume_recording()
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)
    fs_file_path = fs.file_path + '-000002.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000002.mp4'
    assert isfile(us_file_path)
    assert timing_report(us_file_path + '.timing.yml')
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (roi[2], roi[3])
    assert resolution(us_file_path) == (roi[2], roi[3])
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # yet another recording
    recording_duration = 25
    fs.set_full_frame()
    fs.resume_recording()
    us.resume_recording()
    us.set_full_frame()  # should have no effect
    sleep(recording_duration+2)
    fs.pause_recording()
    us.pause_recording()
    sleep(4)
    fs_file_path = fs.file_path + '-000003.mp4'
    assert timing_report(fs_file_path + '.timing.yml')
    assert isfile(fs_file_path)
    us_file_path = us.file_path + '-000003.mp4'
    assert timing_report(us_file_path + '.timing.yml')
    assert isfile(us_file_path)
    assert frame_rate(fs_file_path) == fs.frame_rate
    assert frame_rate(us_file_path) == us.frame_rate
    assert duration(fs_file_path) >= recording_duration
    assert duration(us_file_path) >= recording_duration
    assert resolution(fs_file_path) == (1920, 1080)
    assert resolution(us_file_path) == (roi[2], roi[3])
    assert codec(fs_file_path) == 'hevc'
    assert codec(us_file_path) == 'hevc'

    # finish session
    fs.stop()
    us.stop()

    # join threads
    fs.join(timeout=fs.timeout_limit)
    assert not fs.isAlive()
    us.join(timeout=us.timeout_limit)
    assert not us.isAlive()

    # test dumped Recorder configuration
    dump(fs)
    dumped_fs = parse(join(dirname(fs_file_path), 'config.yml'))
    assert dirname(dirname(dumped_fs.file_path)) == dirname(fs.file_path)
    assert dumped_fs.frame_rate == fs.frame_rate
    assert dumped_fs.port == fs.port
    assert dumped_fs.timeout_limit == fs.timeout_limit
    dump(us)
    dumped_us = parse(join(dirname(us_file_path), 'config.yml'))
    assert dirname(dirname(dumped_us.file_path)) == dirname(us.file_path)
    assert dumped_us.frame_rate == us.frame_rate
    assert dumped_us.port == us.port
    assert dumped_us.timeout_limit == us.timeout_limit

@pytest.yield_fixture(autouse=True)
def cleanup():
    # This section runs before each test

    # Run test
    yield
    # This section runs after each test
    for recorder in recorders:
        recorder.stop()
    sleep(5)

    for recorder in recorders:
        recorder.join(timeout=recorder.timeout_limit)
    sleep(5)

    for recorder in recorders:
        rmtree(dirname(recorder.file_path))

    for recorder in recorders:
        assert not recorder.isAlive()
<|endoftext|>