<commit_before>'''
@date 2015-08-20
@author Hong-She Liang <starofrainnight@gmail.com>
'''

import re
import os
import io
import sys
import math
import urllib
import pip
import pkg_resources
import struct
import platform
import six
from six.moves import urllib
from setuptools.command.install import install as distutils_install
from setuptools.command.easy_install import is_64bit
from pkg_resources._vendor.packaging.version import Version
from ..downloader import download, download_file_insecure_to_io
from ... import windows_api
from pip.wheel import Wheel

class BaseUwbpepServer(object):
    def __init__(self):
        pass

    def _get_page_url(self):
        raise NotImplemented()

    def _get_headers(self):
        return {'User-Agent': 'Mozilla/5.0 (Windows NT 6.3; rv:36.0) Gecko/20100101 Firefox/36.0'}

    def _parse_urls(self, content):
        raise NotImplemented()

    def _decode_url_and_filename(self, amatch):
        raise NotImplemented()

    def _get_index_page(self):
        bytes_io = io.BytesIO()
        try:
            download_file_insecure_to_io(
                self._get_page_url(),
                bytes_io,
                # If you pass a wrong user agent,
                headers=self._get_headers(),
                )
            return bytes_io.getvalue().decode('utf-8')
        finally:
            bytes_io.close()

    def connect(self):
        print('Downloading list page of "Unofficial Windows Binaries for Python Extension Packages" ...')

        content = self._get_index_page()

        print("Download finished. \nParsing ...")

        matches = self._parse_urls(content)

        # Initialize packages with names
        packages = {}

        # Decrypt links
        for amatch in matches:
            url, filename = self._decode_url_and_filename(amatch)

            filebasename, fileextname = os.path.splitext(filename)
            if fileextname not in [".whl", ".zip", ".exe"]:
                continue

            if len(filename.split("-")) < 5:
                continue

            if fileextname == ".exe":
                # Fixed *.exe name to fit for Wheel() requirement! A slight trick
                # to support *.exe package.

                filebasename = filebasename.replace('.win-amd64', '-win_amd64')
                filebasename = filebasename.replace('.win', '-win')
                filebasename = filebasename.replace('-py2.', '-cp2')
                filebasename = filebasename.replace('-py3.', '-cp3')

                wheel_info = filebasename.split('-')
                filename = "%s-%s-%s-%s-%s.whl" % (
                    wheel_info[0],
                    wheel_info[1],
                    wheel_info[3],
                    'none',
                    wheel_info[2]
                    )
            elif fileextname == ".zip":
                filename = "%s%s" % (filebasename, ".whl")

            wheel = Wheel(filename)

            package_name = wheel.name.lower().replace("_", "-")
            if package_name not in packages:
                packages[package_name] = {}
                packages[package_name]["wheels"] = []
                packages[package_name]["requirements"] = []

            packages[package_name]["wheels"].append((wheel, url))

        self.packages = packages

    def find_package(self, requirement_text):
        requirement = pkg_resources.Requirement.parse(requirement_text)
        wheel_contexts = self.packages[requirement.key]["wheels"]

        if is_64bit():
            python_platform = "64"
        else:
            python_platform = "32"

        python_versions = set([
            "cp%s" % platform.python_version_tuple()[0],
            "cp%s%s" % (platform.python_version_tuple()[0], platform.python_version_tuple()[1]),
            "py%s" % platform.python_version_tuple()[0],
            "py%s%s" % (platform.python_version_tuple()[0], platform.python_version_tuple()[1]),
            ])

        for wheel, url in wheel_contexts:
            if python_platform not in wheel.plats[0]:
                continue

            if len(set(wheel.pyversions) & python_versions) <= 0:
                continue

            wheel_version = Version(wheel.version)
            if not requirement.specifier.contains(wheel_version):
                continue

            return (wheel, url)

        raise KeyError("Can't find the requirement : %s" % requirement_text)

    def install(self, requirement):
        wheel, url = self.find_package(requirement)
        filename = wheel.filename

        print("Downloading ... %s " % url)

        download(url, filename, self._get_headers())

        pip.main(["install", filename])

        # Only pywin32 needs postinstall script.
        if filename.startswith("pywin32"):
            postinstall_script_path = os.path.normpath(os.path.join(
                    sys.base_exec_prefix,
                    "Scripts",
                    "pywin32_postinstall.py"))
            postinstall_script_path = postinstall_script_path.replace("/", "\\")
            windows_api.RunAsAdmin([
                sys.executable,
                postinstall_script_path])

class PythonlibsUwbpepServer(BaseUwbpepServer):
    page_url = "http://www.lfd.uci.edu/~gohlke/pythonlibs"

    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _decode_url(self, ml, mi):
        mi = mi.replace('<','<').replace('>','>').replace('&','&')

        ot=""
        for j in range(0, len(mi)):
            ot += chr(ml[ord(mi[j])-48])

        return ot

    def _get_page_url(self):
        return "http://www.lfd.uci.edu/~gohlke/pythonlibs"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'javascript:dl\(\[([^\]]*?)\],\s*"([^"]*?)"', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        ml = []
        for number in amatch[0].split(","):
            ml.append(int(number))

        url = self._decode_url(ml, amatch[1])
        url = "%s/%s" % (self._get_page_url(), url)
        filename = os.path.basename(url)

        return (url, filename)

class GithubUwbpepServer(BaseUwbpepServer):
    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _get_page_url(self):
        return "https://github.com/starofrainnight/uwbpep/releases/tag/v1.0"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        url_parties = list(urllib.parse.urlparse(self._get_page_url()))
        url_parties[2] = amatch
        url = urllib.parse.urlunparse(url_parties)
        filename = os.path.basename(url)

        return (url, filename)

class PypiUwbpepServer(BaseUwbpepServer):
    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _get_page_url(self):
        return "https://pypi.python.org/pypi/uwbpep/1.0"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'<a href="([^"]*?\.whl#md5=[^"]*?)"', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        url = amatch
        filename = os.path.basename(url.split("#")[0])
        # Removed the first "uwbpep1.0_" tag!
        filename = filename[len("uwbpep1.0_"):]

        return (url, filename)

class install(distutils_install):
    """
    The install command provide extension packages automatic install from
    our custom package page.

    We make a cached page at

    https://github.com/starofrainnight/uwbpep/releases/tag/v1.0

    which copied some needed packages from

    http://www.lfd.uci.edu/~gohlke/pythonlibs/

    Abbreviations:

    uwbpep : Unofficial Windows Binaries for Python Extension Packages

    """
    def _prepare_requirements(self):
        # Try to use pip install first
        failed_requires = ["pywin32"]
        for arequire in self.distribution.install_requires:

            # pip.main() will return 0 while successed ..
            if pip.main(["install", arequire]) != 0:
                failed_requires.append(arequire)

        if len(failed_requires) > 0:
            server = PypiUwbpepServer()
            server.connect()

            # Try to install failed requires from UWBPEP
            for arequire in failed_requires:
                server.install(arequire)

    def run(self):
        if sys.platform == "win32":
            self._prepare_requirements()

        # self.distribution.install_requires
        distutils_install.run(self)

        <commit_msg>Renamed find_package()  to _find_package() for hide it from others<commit_after>'''
@date 2015-08-20
@author Hong-She Liang <starofrainnight@gmail.com>
'''

import re
import os
import io
import sys
import math
import urllib
import pip
import pkg_resources
import struct
import platform
import six
from six.moves import urllib
from setuptools.command.install import install as distutils_install
from setuptools.command.easy_install import is_64bit
from pkg_resources._vendor.packaging.version import Version
from ..downloader import download, download_file_insecure_to_io
from ... import windows_api
from pip.wheel import Wheel

class BaseUwbpepServer(object):
    def __init__(self):
        pass

    def _get_page_url(self):
        raise NotImplemented()

    def _get_headers(self):
        return {'User-Agent': 'Mozilla/5.0 (Windows NT 6.3; rv:36.0) Gecko/20100101 Firefox/36.0'}

    def _parse_urls(self, content):
        raise NotImplemented()

    def _decode_url_and_filename(self, amatch):
        raise NotImplemented()

    def _get_index_page(self):
        bytes_io = io.BytesIO()
        try:
            download_file_insecure_to_io(
                self._get_page_url(),
                bytes_io,
                # If you pass a wrong user agent,
                headers=self._get_headers(),
                )
            return bytes_io.getvalue().decode('utf-8')
        finally:
            bytes_io.close()

    def connect(self):
        print('Downloading list page of "Unofficial Windows Binaries for Python Extension Packages" ...')

        content = self._get_index_page()

        print("Download finished. \nParsing ...")

        matches = self._parse_urls(content)

        # Initialize packages with names
        packages = {}

        # Decrypt links
        for amatch in matches:
            url, filename = self._decode_url_and_filename(amatch)

            filebasename, fileextname = os.path.splitext(filename)
            if fileextname not in [".whl", ".zip", ".exe"]:
                continue

            if len(filename.split("-")) < 5:
                continue

            if fileextname == ".exe":
                # Fixed *.exe name to fit for Wheel() requirement! A slight trick
                # to support *.exe package.

                filebasename = filebasename.replace('.win-amd64', '-win_amd64')
                filebasename = filebasename.replace('.win', '-win')
                filebasename = filebasename.replace('-py2.', '-cp2')
                filebasename = filebasename.replace('-py3.', '-cp3')

                wheel_info = filebasename.split('-')
                filename = "%s-%s-%s-%s-%s.whl" % (
                    wheel_info[0],
                    wheel_info[1],
                    wheel_info[3],
                    'none',
                    wheel_info[2]
                    )
            elif fileextname == ".zip":
                filename = "%s%s" % (filebasename, ".whl")

            wheel = Wheel(filename)

            package_name = wheel.name.lower().replace("_", "-")
            if package_name not in packages:
                packages[package_name] = {}
                packages[package_name]["wheels"] = []
                packages[package_name]["requirements"] = []

            packages[package_name]["wheels"].append((wheel, url))

        self.packages = packages

    def _find_package(self, requirement_text):
        requirement = pkg_resources.Requirement.parse(requirement_text)
        wheel_contexts = self.packages[requirement.key]["wheels"]

        if is_64bit():
            python_platform = "64"
        else:
            python_platform = "32"

        python_versions = set([
            "cp%s" % platform.python_version_tuple()[0],
            "cp%s%s" % (platform.python_version_tuple()[0], platform.python_version_tuple()[1]),
            "py%s" % platform.python_version_tuple()[0],
            "py%s%s" % (platform.python_version_tuple()[0], platform.python_version_tuple()[1]),
            ])

        for wheel, url in wheel_contexts:
            if python_platform not in wheel.plats[0]:
                continue

            if len(set(wheel.pyversions) & python_versions) <= 0:
                continue

            wheel_version = Version(wheel.version)
            if not requirement.specifier.contains(wheel_version):
                continue

            return (wheel, url)

        raise KeyError("Can't find the requirement : %s" % requirement_text)

    def install(self, requirement):
        wheel, url = self._find_package(requirement)
        filename = wheel.filename

        print("Downloading ... %s " % url)

        download(url, filename, self._get_headers())

        pip.main(["install", filename])

        # Only pywin32 needs postinstall script.
        if filename.startswith("pywin32"):
            postinstall_script_path = os.path.normpath(os.path.join(
                    sys.base_exec_prefix,
                    "Scripts",
                    "pywin32_postinstall.py"))
            postinstall_script_path = postinstall_script_path.replace("/", "\\")
            windows_api.RunAsAdmin([
                sys.executable,
                postinstall_script_path])

class PythonlibsUwbpepServer(BaseUwbpepServer):
    page_url = "http://www.lfd.uci.edu/~gohlke/pythonlibs"

    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _decode_url(self, ml, mi):
        mi = mi.replace('<','<').replace('>','>').replace('&','&')

        ot=""
        for j in range(0, len(mi)):
            ot += chr(ml[ord(mi[j])-48])

        return ot

    def _get_page_url(self):
        return "http://www.lfd.uci.edu/~gohlke/pythonlibs"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'javascript:dl\(\[([^\]]*?)\],\s*"([^"]*?)"', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        ml = []
        for number in amatch[0].split(","):
            ml.append(int(number))

        url = self._decode_url(ml, amatch[1])
        url = "%s/%s" % (self._get_page_url(), url)
        filename = os.path.basename(url)

        return (url, filename)

class GithubUwbpepServer(BaseUwbpepServer):
    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _get_page_url(self):
        return "https://github.com/starofrainnight/uwbpep/releases/tag/v1.0"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        url_parties = list(urllib.parse.urlparse(self._get_page_url()))
        url_parties[2] = amatch
        url = urllib.parse.urlunparse(url_parties)
        filename = os.path.basename(url)

        return (url, filename)

class PypiUwbpepServer(BaseUwbpepServer):
    def __init__(self):
        super(BaseUwbpepServer, self).__init__()

    def _get_page_url(self):
        return "https://pypi.python.org/pypi/uwbpep/1.0"

    def _parse_urls(self, content):
        re_flags = re.DOTALL|re.MULTILINE
        return re.findall(r'<a href="([^"]*?\.whl#md5=[^"]*?)"', content, re_flags)

    def _decode_url_and_filename(self, amatch):
        url = amatch
        filename = os.path.basename(url.split("#")[0])
        # Removed the first "uwbpep1.0_" tag!
        filename = filename[len("uwbpep1.0_"):]

        return (url, filename)

class install(distutils_install):
    """
    The install command provide extension packages automatic install from
    our custom package page.

    We make a cached page at

    https://github.com/starofrainnight/uwbpep/releases/tag/v1.0

    which copied some needed packages from

    http://www.lfd.uci.edu/~gohlke/pythonlibs/

    Abbreviations:

    uwbpep : Unofficial Windows Binaries for Python Extension Packages

    """
    def _prepare_requirements(self):
        # Try to use pip install first
        failed_requires = ["pywin32"]
        for arequire in self.distribution.install_requires:

            # pip.main() will return 0 while successed ..
            if pip.main(["install", arequire]) != 0:
                failed_requires.append(arequire)

        if len(failed_requires) > 0:
            server = PypiUwbpepServer()
            server.connect()

            # Try to install failed requires from UWBPEP
            for arequire in failed_requires:
                server.install(arequire)

    def run(self):
        if sys.platform == "win32":
            self._prepare_requirements()

        # self.distribution.install_requires
        distutils_install.run(self)

        <|endoftext|>


<commit_before>from unittest import TestCase

import jsonpickle

from cloudshell.cp.azure.common.deploy_data_holder import DeployDataHolder

from tests.helpers.test_helper import TestHelper

from cloudshell.cp.azure.domain.services.key_pair import KeyPairService

from cloudshell.cp.azure.domain.services.tags import TagService

from cloudshell.cp.azure.domain.services.network_service import NetworkService

from cloudshell.cp.azure.domain.services.virtual_machine_service import VirtualMachineService

from cloudshell.cp.azure.domain.services.storage_service import StorageService

from cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation import \
    PrepareConnectivityOperation
from mock import MagicMock, Mock

class TestPrepareConnectivity(TestCase):
    def setUp(self):
        self.storage_service = StorageService()
        self.vm_service = VirtualMachineService()
        self.network_service = NetworkService()
        self.tag_service = TagService()
        self.key_pair_service = KeyPairService()
        self.logger = MagicMock()

        self.prepare_connectivity_operation = PrepareConnectivityOperation(logger=self.logger,
                                                                           vm_service=self.vm_service,
                                                                           network_service=self.network_service,
                                                                           storage_service=self.storage_service,
                                                                           tags_service=self.tag_service,
                                                                           key_pair_service=self.key_pair_service)

    def test_prepare_connectivity(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.network_service.get_virtual_network_by_tag = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = '10.0.0.0/12'
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.prepare_connectivity_operation.prepare_connectivity(
            reservation=MagicMock(),
            cloud_provider_model=MagicMock(),
            storage_client=MagicMock(),
            resource_client=MagicMock(),
            network_client=MagicMock(),
            logger=self.logger,
            request=prepare_connectivity_request)

        # Verify
        # Created resource Group
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.vm_service.create_resource_group))

        # created Storage account
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.storage_service.create_storage_account))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_network_by_tag, 2))

        # key pair created
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.key_pair_service.save_key_pair))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_networks))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.create_subnet))

    def test_extract_cidr_throws_error(self):
        action = MagicMock()

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

        action = Mock()
        att = Mock()
        att.attributeName = 'Network'
        att.attributeValue = [1,2,3]
        action.customActionAttributes = [att]

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

    def test_prepare_connectivity_throes_exception_on_unavailable_VNETs(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()
        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [None, Mock()]

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = '10.0.0.0/12'
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)

        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [Mock(), None]

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)
<commit_msg>fix tests<commit_after>from unittest import TestCase

import jsonpickle

from cloudshell.cp.azure.common.deploy_data_holder import DeployDataHolder

from tests.helpers.test_helper import TestHelper

from cloudshell.cp.azure.domain.services.key_pair import KeyPairService

from cloudshell.cp.azure.domain.services.tags import TagService

from cloudshell.cp.azure.domain.services.network_service import NetworkService

from cloudshell.cp.azure.domain.services.virtual_machine_service import VirtualMachineService

from cloudshell.cp.azure.domain.services.storage_service import StorageService

from cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation import \
    PrepareConnectivityOperation
from mock import MagicMock, Mock

class TestPrepareConnectivity(TestCase):
    def setUp(self):
        self.storage_service = StorageService()
        self.vm_service = VirtualMachineService()
        self.network_service = NetworkService()
        self.tag_service = TagService()
        self.key_pair_service = KeyPairService()
        self.logger = MagicMock()

        self.prepare_connectivity_operation = PrepareConnectivityOperation(logger=self.logger,
                                                                           vm_service=self.vm_service,
                                                                           network_service=self.network_service,
                                                                           storage_service=self.storage_service,
                                                                           tags_service=self.tag_service,
                                                                           key_pair_service=self.key_pair_service)

    def test_prepare_connectivity(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.network_service.get_virtual_network_by_tag = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = [1]
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.prepare_connectivity_operation.prepare_connectivity(
            reservation=MagicMock(),
            cloud_provider_model=MagicMock(),
            storage_client=MagicMock(),
            resource_client=MagicMock(),
            network_client=MagicMock(),
            logger=self.logger,
            request=prepare_connectivity_request)

        # Verify
        # Created resource Group
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.vm_service.create_resource_group))

        # created Storage account
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.storage_service.create_storage_account))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_network_by_tag, 2))

        # key pair created
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.key_pair_service.save_key_pair))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_networks))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.create_subnet))

    def test_extract_cidr_throws_error(self):
        action = MagicMock()

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

        action = Mock()
        att = Mock()
        att.attributeName = 'Network'
        att.attributeValue = [1,2,3]
        action.customActionAttributes = [att]

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

    def test_prepare_connectivity_throes_exception_on_unavailable_VNETs(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()
        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [None, Mock()]

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = '10.0.0.0/12'
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)

        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [Mock(), None]

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)
<|endoftext|>


<commit_before>import cloudpickle
import fairing
import glob
import os
import sys
import types
import tempfile
from pathlib import Path
from enum import Enum

from fairing.constants import constants
from .base import BasePreProcessor
from fairing.functions.function_shim import get_execution_obj_type, ObjectType

FUNCTION_SHIM = 'function_shim.py'
SERIALIZED_FN_FILE = 'pickled_fn.p'

# TODO(@karthikv2k): Ref #122 Find a better way to support deployer specific preprocessing
OUTPUT_FILE = """import cloudpickle
{OBJ_NAME} = cloudpickle.load(open("{SERIALIZED_FN_FILE}", "rb"))
"""

class FunctionPreProcessor(BasePreProcessor):
    """
    FunctionPreProcessor preprocesses a single function.
    It sets as the command a function_shim that calls the function directly.

    args: function_name - the name of the function to be called
    """
    def __init__(self,
                 function_obj,
                 path_prefix=constants.DEFAULT_DEST_PREFIX,
                 output_map=None,
                 input_files=None):
        super().__init__(
            output_map=output_map,
            path_prefix=path_prefix,
            input_files=input_files)

        if get_execution_obj_type(function_obj) ==  ObjectType.NOT_SUPPORTED:
            raise RuntimeError("Object must of type function or a class")

        fairing_dir = os.path.dirname(fairing.__file__)
        self.output_map[os.path.join(fairing_dir, "functions", FUNCTION_SHIM)] = \
            os.path.join(path_prefix, FUNCTION_SHIM)

        # Make sure the user code can be imported as a module
        self.output_map[os.path.join(fairing_dir, '__init__.py')] = \
            os.path.join(path_prefix, '__init__.py')

        # Make sure fairing can use imported as a module
        self.output_map[os.path.join(fairing_dir, '__init__.py')] = \
            os.path.join(path_prefix, "fairing", '__init__.py')

        # Make sure cloudpickle can be imported as a module
        cloudpickle_dir = os.path.dirname(cloudpickle.__file__)
        self.output_map[os.path.join(cloudpickle_dir, '__init__.py')] = \
            os.path.join(path_prefix, "cloudpickle", '__init__.py')
        self.output_map[os.path.join(cloudpickle_dir, 'cloudpickle.py')] = \
            os.path.join(path_prefix, "cloudpickle", 'cloudpickle.py')

        _, temp_payload_file = tempfile.mkstemp()
        with open(temp_payload_file, "wb") as f:
            cloudpickle.dump(function_obj, f)
        # Adding the serialized file to the context
        payload_file_in_context = os.path.join(path_prefix, SERIALIZED_FN_FILE)
        self.output_map[temp_payload_file] = payload_file_in_context

        # TODO(@karthikv2k): Ref #122 Find a better way to support deployer specific preprocessing
        _, temp_payload_wrapper_file = tempfile.mkstemp()
        with open(temp_payload_wrapper_file, "w") as f:
            contents  = OUTPUT_FILE.format(OBJ_NAME=function_obj.__name__, SERIALIZED_FN_FILE=SERIALIZED_FN_FILE)
            f.write(contents)
        # Adding the serialized file to the context
        payload_wrapper_file_in_context = os.path.join(path_prefix, function_obj.__name__ + ".py")
        self.output_map[temp_payload_wrapper_file] = payload_wrapper_file_in_context

        local_python_version = ".".join([str(x) for x in sys.version_info[0:3]])

        self.command = ["python", os.path.join(self.path_prefix, FUNCTION_SHIM),
                        "--serialized_fn_file", payload_file_in_context,
                        "--python_version", local_python_version]

    def get_command(self):
        return self.command
<commit_msg>Warn users who are using function preprocessor and not in notebook (#292)<commit_after>import cloudpickle
import fairing
import glob
import logging
import os
import sys
import types
import tempfile
from pathlib import Path
from enum import Enum

from fairing.constants import constants
from .base import BasePreProcessor
from fairing.functions.function_shim import get_execution_obj_type, ObjectType
from fairing.notebook import notebook_util

logger = logging.getLogger(__name__)

FUNCTION_SHIM = 'function_shim.py'
SERIALIZED_FN_FILE = 'pickled_fn.p'

# TODO(@karthikv2k): Ref #122 Find a better way to support deployer specific preprocessing
OUTPUT_FILE = """import cloudpickle
{OBJ_NAME} = cloudpickle.load(open("{SERIALIZED_FN_FILE}", "rb"))
"""

class FunctionPreProcessor(BasePreProcessor):
    """
    FunctionPreProcessor preprocesses a single function.
    It sets as the command a function_shim that calls the function directly.

    args: function_name - the name of the function to be called
    """
    def __init__(self,
                 function_obj,
                 path_prefix=constants.DEFAULT_DEST_PREFIX,
                 output_map=None,
                 input_files=None):
        super().__init__(
            output_map=output_map,
            path_prefix=path_prefix,
            input_files=input_files)

        if not notebook_util.is_in_notebook():
            logger.warning("The FunctionPreProcessor is optimized for using in a notebook or IPython environment. "
                           "For it to work, the python version should be same for both local python and the python in "
                           "the docker. Please look at alternatives like BasePreprocessor or FullNotebookPreprocessor.")

        if get_execution_obj_type(function_obj) ==  ObjectType.NOT_SUPPORTED:
            raise RuntimeError("Object must of type function or a class")

        fairing_dir = os.path.dirname(fairing.__file__)
        self.output_map[os.path.join(fairing_dir, "functions", FUNCTION_SHIM)] = \
            os.path.join(path_prefix, FUNCTION_SHIM)

        # Make sure the user code can be imported as a module
        self.output_map[os.path.join(fairing_dir, '__init__.py')] = \
            os.path.join(path_prefix, '__init__.py')

        # Make sure fairing can use imported as a module
        self.output_map[os.path.join(fairing_dir, '__init__.py')] = \
            os.path.join(path_prefix, "fairing", '__init__.py')

        # Make sure cloudpickle can be imported as a module
        cloudpickle_dir = os.path.dirname(cloudpickle.__file__)
        self.output_map[os.path.join(cloudpickle_dir, '__init__.py')] = \
            os.path.join(path_prefix, "cloudpickle", '__init__.py')
        self.output_map[os.path.join(cloudpickle_dir, 'cloudpickle.py')] = \
            os.path.join(path_prefix, "cloudpickle", 'cloudpickle.py')

        _, temp_payload_file = tempfile.mkstemp()
        with open(temp_payload_file, "wb") as f:
            cloudpickle.dump(function_obj, f)
        # Adding the serialized file to the context
        payload_file_in_context = os.path.join(path_prefix, SERIALIZED_FN_FILE)
        self.output_map[temp_payload_file] = payload_file_in_context

        # TODO(@karthikv2k): Ref #122 Find a better way to support deployer specific preprocessing
        _, temp_payload_wrapper_file = tempfile.mkstemp()
        with open(temp_payload_wrapper_file, "w") as f:
            contents  = OUTPUT_FILE.format(OBJ_NAME=function_obj.__name__, SERIALIZED_FN_FILE=SERIALIZED_FN_FILE)
            f.write(contents)
        # Adding the serialized file to the context
        payload_wrapper_file_in_context = os.path.join(path_prefix, function_obj.__name__ + ".py")
        self.output_map[temp_payload_wrapper_file] = payload_wrapper_file_in_context

        local_python_version = ".".join([str(x) for x in sys.version_info[0:3]])

        self.command = ["python", os.path.join(self.path_prefix, FUNCTION_SHIM),
                        "--serialized_fn_file", payload_file_in_context,
                        "--python_version", local_python_version]

    def get_command(self):
        return self.command
<|endoftext|>


<commit_before>import os
import tempfile

from django.test import TestCase
from viewflow.graphviz import diagram
from viewflow.test import FlowTest

from examples.shipment.flows import ShipmentFlow
from examples.shipment.models import Carrier

class ShipmentFlowDiagrammTests(TestCase):
    def test_diagram_creation_succeed(self):
        outdir = tempfile.mkdtemp()
        png_filename = os.path.join(outdir, 'output.png')

        diagram(ShipmentFlow, png_filename)

        self.assertTrue(os.path.exists(png_filename))

class ShipmentFlowTests(TestCase):
    def setUp(self):
        self.default_carrier = Carrier.objects.get(name=Carrier.DEFAULT)
        self.special_carrier = Carrier.objects.exclude(name=Carrier.DEFAULT)[0]

    def _test_normal_post_succeed(self):
        with FlowTest(ShipmentFlow) as flow:
            # Clerk start process
            flow.Task(ShipmentFlow.start).User('shipment/clerk') \
                .Execute({'goods_tag': 'TST123'}) \
                .Assert(lambda p: p.created is not None)

            # Clerk decides that is special shipment post
            flow.Task(ShipmentFlow.shipment_type).User('shipment/clerk') \
                .Execute({'carrier': self.special_carrier.pk}) \
                .Assert(lambda p: not p.is_normal_post())

            # Clerk inputs carrier quote
            flow.Task(ShipmentFlow.request_quotes).User('shipment/clerk') \
                .Execute({'carrier_quote': 1000}) \
                .Assert(lambda p: p.shipment.carrier_quote == 1000)

            # Worker package goods
            flow.Task(ShipmentFlow.package_goods).User('shipment/worker') \
                .Execute()

            # Worker added paperwork and moved package out
            flow.Task(ShipmentFlow.move_package).User('shipment/worker') \
                .Execute() \
                .Assert(lambda p: p.finished is not None)

    def test_insured_post_succeed(self):
        with FlowTest(ShipmentFlow) as flow:
            # Clerk start process
            flow.Task(ShipmentFlow.start).User('shipment/clerk') \
                .Execute({'goods_tag': 'TST123'}) \
                .Assert(lambda p: p.created is not None)

            # Clerk decides that is normal shipment post
            flow.Task(ShipmentFlow.shipment_type).User('shipment/clerk') \
                .Execute({'carrier': self.default_carrier.pk}) \
                .Assert(lambda p: p.is_normal_post())

            # Clerk decides that extra insurance required
            flow.Task(ShipmentFlow.check_insurance).User('shipment/clerk') \
                .Execute({'need_insurance': True}) \
                .Assert(lambda p: p.need_extra_insurance())

            # Clerk fills post label
            flow.Task(ShipmentFlow.fill_post_label).User('shipment/clerk') \
                .Execute({'post_label': 'Test label'}) \
                .Assert(lambda p: p.shipment.post_label == 'Test label')

            # Worker package goods
            flow.Task(ShipmentFlow.package_goods).User('shipment/worker') \
                .Execute()

            # Manager takes extra insurance
            flow.Task(ShipmentFlow.take_extra_insurance).User('shipment/manager') \
                .Execute({'company_name': 'Test company',
                          'cost': 100}) \
                .Assert(lambda p: p.shipment.insurance.company_name == 'Test company')

            # Worker added paperwork and moved package out
            flow.Task(ShipmentFlow.move_package).User('shipment/worker') \
                .Execute() \
                .Assert(lambda p: p.finished is not None)
<commit_msg>Enable shipment tests<commit_after>import os
import tempfile

from django.test import TestCase
from viewflow.graphviz import diagram
from viewflow.test import FlowTest

from examples.shipment.flows import ShipmentFlow
from examples.shipment.models import Carrier

class ShipmentFlowDiagrammTests(TestCase):
    def test_diagram_creation_succeed(self):
        outdir = tempfile.mkdtemp()
        png_filename = os.path.join(outdir, 'output.png')

        diagram(ShipmentFlow, png_filename)

        self.assertTrue(os.path.exists(png_filename))

class ShipmentFlowTests(TestCase):
    def setUp(self):
        self.default_carrier = Carrier.objects.get(name=Carrier.DEFAULT)
        self.special_carrier = Carrier.objects.exclude(name=Carrier.DEFAULT)[0]

    def test_normal_post_succeed(self):
        with FlowTest(ShipmentFlow) as flow:
            # Clerk start process
            flow.Task(ShipmentFlow.start).User('shipment/clerk') \
                .Execute({'goods_tag': 'TST123'}) \
                .Assert(lambda p: p.created is not None)

            # Clerk decides that is special shipment post
            flow.Task(ShipmentFlow.shipment_type).User('shipment/clerk') \
                .Execute({'carrier': self.special_carrier.pk}) \
                .Assert(lambda p: not p.is_normal_post())

            # Clerk inputs carrier quote
            flow.Task(ShipmentFlow.request_quotes).User('shipment/clerk') \
                .Execute({'carrier_quote': 1000}) \
                .Assert(lambda p: p.shipment.carrier_quote == 1000)

            # Worker package goods
            flow.Task(ShipmentFlow.package_goods).User('shipment/worker') \
                .Execute()

            # Worker added paperwork and moved package out
            flow.Task(ShipmentFlow.move_package).User('shipment/worker') \
                .Execute() \
                .Assert(lambda p: p.finished is not None)

    def test_insured_post_succeed(self):
        with FlowTest(ShipmentFlow) as flow:
            # Clerk start process
            flow.Task(ShipmentFlow.start).User('shipment/clerk') \
                .Execute({'goods_tag': 'TST123'}) \
                .Assert(lambda p: p.created is not None)

            # Clerk decides that is normal shipment post
            flow.Task(ShipmentFlow.shipment_type).User('shipment/clerk') \
                .Execute({'carrier': self.default_carrier.pk}) \
                .Assert(lambda p: p.is_normal_post())

            # Clerk decides that extra insurance required
            flow.Task(ShipmentFlow.check_insurance).User('shipment/clerk') \
                .Execute({'need_insurance': True}) \
                .Assert(lambda p: p.need_extra_insurance())

            # Clerk fills post label
            flow.Task(ShipmentFlow.fill_post_label).User('shipment/clerk') \
                .Execute({'post_label': 'Test label'}) \
                .Assert(lambda p: p.shipment.post_label == 'Test label')

            # Worker package goods
            flow.Task(ShipmentFlow.package_goods).User('shipment/worker') \
                .Execute()

            # Manager takes extra insurance
            flow.Task(ShipmentFlow.take_extra_insurance).User('shipment/manager') \
                .Execute({'company_name': 'Test company',
                          'cost': 100}) \
                .Assert(lambda p: p.shipment.insurance.company_name == 'Test company')

            # Worker added paperwork and moved package out
            flow.Task(ShipmentFlow.move_package).User('shipment/worker') \
                .Execute() \
                .Assert(lambda p: p.finished is not None)
<|endoftext|>


<commit_before>
import time
import os

from devops.model import Environment, Node, Network, Interface
from devops.controller import Controller
from devops.driver.libvirt import Libvirt

from devops import yaml_config_loader

INSTALLATION_ISO = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'install.iso'))
INSTALLATION_ISO_URL = "http://mc0n1-srt.srt.mirantis.net/nailgun-ubuntu-12.04-amd64.last.iso"

MASTER_AND_SLAVE_CONFIG = """
name: 'Sample environment'
networks:
  - network: internal
  - network: external
    type: bridged
nodes:
  - node: gateway
    disk: '5Gb'
    cdrom: '%s'
    networks: ['external', 'internal']
  - node: slug
    networks: ['internal']
""" % (INSTALLATION_ISO,)

def download_iso():
    if not os.path.exists(INSTALLATION_ISO):
        print("Downloading installation iso")
        os.system("wget -O '%s' -c '%s'" % (INSTALLATION_ISO, INSTALLATION_ISO_URL))
        print("Finished downloading installation iso")

def main():
    download_iso()

    environment = yaml_config_loader.load(MASTER_AND_SLAVE_CONFIG)

    external_network = environment.network['external']

    gateway = environment.node['gateway']
    gateway.vnc = True

    print("Creating environment")

    controller = Controller(Libvirt())
    controller.build_environment(environment)

    print("Environment created")

    print("Starting node")
    gateway.start()

    print("VNC is available on vnc://localhost:%d" % gateway.vnc_port)

    print("Waiting node to boot")
    time.sleep(10)
    print("Sending user input")
    gateway.send_keys("""<Esc><Enter>
<Wait>
/install/vmlinuz initrd=/install/initrd.gz
 priority=critical
 locale=en_US
 file=/cdrom/preseed/manual.seed
 vga=788
 netcfg/get_ipaddress=10.12.0.2
 netcfg/get_netmask=255.255.255.0
 netcfg/get_gateway=10.12.0.1
 netcfg/get_nameservers=10.12.0.1
 netcfg/confirm_static=true
 <Enter>""")
    print("Finished sending user input")

if __name__ == '__main__':
    main()

<commit_msg>[devops] Updated run.py: pass network settings to VM kernel arguments<commit_after>
import time
import os

from devops.model import Environment, Node, Network, Interface
from devops.controller import Controller
from devops.driver.libvirt import Libvirt

from devops import yaml_config_loader

INSTALLATION_ISO = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'install.iso'))
INSTALLATION_ISO_URL = "http://mc0n1-srt.srt.mirantis.net/nailgun-ubuntu-12.04-amd64.last.iso"

MASTER_AND_SLAVE_CONFIG = """
name: 'Sample environment'
networks:
  - network: internal
  - network: external
    type: bridged
nodes:
  - node: gateway
    disk: '5Gb'
    cdrom: '%s'
    networks: ['external', 'internal']
  - node: slug
    networks: ['internal']
""" % (INSTALLATION_ISO,)

def download_iso():
    if not os.path.exists(INSTALLATION_ISO):
        print("Downloading installation iso")
        os.system("wget -O '%s' -c '%s'" % (INSTALLATION_ISO, INSTALLATION_ISO_URL))
        print("Finished downloading installation iso")

def main():
    download_iso()

    environment = yaml_config_loader.load(MASTER_AND_SLAVE_CONFIG)

    external_network = environment.network['external']

    gateway = environment.node['gateway']
    gateway.vnc = True

    print("Creating environment")

    controller = Controller(Libvirt())
    controller.build_environment(environment)

    print("Environment created")

    print("Starting node")
    gateway.start()

    print("VNC is available on vnc://localhost:%d" % gateway.vnc_port)

    print("Waiting node to boot")
    time.sleep(15)

    print("Sending user input")

    ip = external_network.ip_addresses

    gateway.send_keys("""<Esc><Enter>
<Wait>
/install/vmlinuz initrd=/install/initrd.gz
 priority=critical
 locale=en_US
 file=/cdrom/preseed/manual.seed
 vga=788
 netcfg/get_ipaddress=%s
 netcfg/get_netmask=%s
 netcfg/get_gateway=%s
 netcfg/get_nameservers=%s
 netcfg/confirm_static=true
 <Enter>""" % (ip[2], ip.netmask, ip[1], ip[1]))
    print("Finished sending user input")

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before>device_parents = {}
try:
    from . import wurfl_data
except ImportError:
    pass
else:

    def get_parents(device):
        if device == 'root' or device is None:
            return []
        device = unicode(device)
        try:
            return device_parents[device]
        except KeyError:
            print device, type(device)
            device_parents[device] = [device] + get_parents(wurfl_data.devices.select_id(device).fall_back)
            return device_parents[device]

    for device in wurfl_data.devices:
        if not device in device_parents:
            device_parents[device] = get_parents(device)

    # Convert our lists to frozensets for efficiency's sake
    for key in device_parents:
        device_parents[key] = frozenset(device_parents[key])
<commit_msg>Removed debugging print statement.<commit_after>device_parents = {}
try:
    from . import wurfl_data
except ImportError:
    pass
else:

    def get_parents(device):
        if device == 'root' or device is None:
            return []
        device = unicode(device)
        try:
            return device_parents[device]
        except KeyError:
            device_parents[device] = [device] + get_parents(wurfl_data.devices.select_id(device).fall_back)
            return device_parents[device]

    for device in wurfl_data.devices:
        if not device in device_parents:
            device_parents[device] = get_parents(device)

    # Convert our lists to frozensets for efficiency's sake
    for key in device_parents:
        device_parents[key] = frozenset(device_parents[key])
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
# File: common.py
# Author: Yuxin Wu <ppwwyyxx@gmail.com>

import tensorflow as tf
import inspect
from functools import wraps
import six
import re
import copy

from ..tfutils.argscope import get_arg_scope
from ..tfutils.model_utils import get_shape_str
from ..utils import logger

# make sure each layer is only logged once
_LAYER_LOGGED = set()
_LAYER_REGISTRY = {}

__all__ = ['layer_register']

class VariableHolder(object):
    """ A proxy to access variables defined in a layer. """
    def __init__(self, **kwargs):
        """
        Args:
            kwargs: {name:variable}
        """
        self._vars = {}
        for k, v in six.iteritems(kwargs):
            self._add_variable(k, v)

    def _add_variable(self, name, var):
        assert name not in self._vars
        self._vars[name] = var

    def __setattr__(self, name, var):
        if not name.startswith('_'):
            self._add_variable(name, var)
        else:
            # private attributes
            super(VariableHolder, self).__setattr__(name, var)

    def __getattr__(self, name):
        return self._vars[name]

    def all(self):
        """
        Returns:
            list of all variables
        """
        return list(six.itervalues(self._vars))

def _register(name, func):
    if name in _LAYER_REGISTRY:
        raise ValueError("Layer named {} is already registered!".format(name))
    if name in ['tf']:
        raise ValueError(logger.error("A layer cannot be named {}".format(name)))
    _LAYER_REGISTRY[name] = func

def get_registered_layer(name):
    """
    Args:
        name (str): the name of the layer, e.g. 'Conv2D'
    Returns:
        the wrapped layer function, or None if not registered.
    """
    return _LAYER_REGISTRY.get(name, None)

def disable_layer_logging():
    """
    Disable the shape logging for all layers from this moment on. Can be
    useful when creating multiple towers.
    """
    class ContainEverything:
        def __contains__(self, x):
            return True
    # can use nonlocal in python3, but how
    globals()['_LAYER_LOGGED'] = ContainEverything()

def layer_register(
        log_shape=False,
        use_scope=True):
    """
    Register a layer.

    Args:
        log_shape (bool): log input/output shape of this layer
        use_scope (bool or None):
            Whether to call this layer with an extra first argument as scope.
            When set to None, it can be called either with or without
            the scope name argument.
            It will try to figure out by checking if the first argument
            is string or not.
    """

    def wrapper(func):
        @wraps(func)
        def wrapped_func(*args, **kwargs):
            assert args[0] is not None, args
            if use_scope:
                name, inputs = args[0], args[1]
                args = args[1:]  # actual positional args used to call func
                assert isinstance(name, six.string_types), name
            else:
                assert not log_shape
                if isinstance(args[0], six.string_types):
                    if use_scope is False:
                        logger.warn(
                            "Please call layer {} without the first scope name argument, "
                            "or register the layer with use_scope=None to allow "
                            "two calling methods.".format(func.__name__))
                    name, inputs = args[0], args[1]
                    args = args[1:]  # actual positional args used to call func
                else:
                    inputs = args[0]
                    name = None
            if not (isinstance(inputs, (tf.Tensor, tf.Variable)) or
                    (isinstance(inputs, (list, tuple)) and
                        isinstance(inputs[0], (tf.Tensor, tf.Variable)))):
                raise ValueError("Invalid inputs to layer: " + str(inputs))

            # use kwargs from current argument scope
            actual_args = copy.copy(get_arg_scope()[func.__name__])
            # explicit kwargs overwrite argscope
            actual_args.update(kwargs)
            # explicit positional args also override argscope
            if six.PY2:
                posargmap = inspect.getcallargs(func, *args)
            else:
                posargmap = inspect.signature(func).bind_partial(*args).arguments
            for k in six.iterkeys(posargmap):
                if k in actual_args:
                    del actual_args[k]

            if name is not None:        # use scope
                with tf.variable_scope(name) as scope:
                    # this name is only used to surpress logging, doesn't hurt to do some heuristics
                    scope_name = re.sub('tower[0-9]+/', '', scope.name)
                    do_log_shape = log_shape and scope_name not in _LAYER_LOGGED
                    if do_log_shape:
                        logger.info("{} input: {}".format(scope.name, get_shape_str(inputs)))

                    # run the actual function
                    outputs = func(*args, **actual_args)

                    if do_log_shape:
                        # log shape info and add activation
                        logger.info("{} output: {}".format(
                            scope.name, get_shape_str(outputs)))
                        _LAYER_LOGGED.add(scope_name)
            else:
                # run the actual function
                outputs = func(*args, **actual_args)
            return outputs

        wrapped_func.symbolic_function = func   # attribute to access the underlying function object
        wrapped_func.use_scope = use_scope
        _register(func.__name__, wrapped_func)
        return wrapped_func

    return wrapper
<commit_msg>fix build for python2<commit_after># -*- coding: UTF-8 -*-
# File: common.py
# Author: Yuxin Wu <ppwwyyxx@gmail.com>

import tensorflow as tf
import inspect
from functools import wraps
import six
import re
import copy

from ..tfutils.argscope import get_arg_scope
from ..tfutils.model_utils import get_shape_str
from ..utils import logger

# make sure each layer is only logged once
_LAYER_LOGGED = set()
_LAYER_REGISTRY = {}

__all__ = ['layer_register']

class VariableHolder(object):
    """ A proxy to access variables defined in a layer. """
    def __init__(self, **kwargs):
        """
        Args:
            kwargs: {name:variable}
        """
        self._vars = {}
        for k, v in six.iteritems(kwargs):
            self._add_variable(k, v)

    def _add_variable(self, name, var):
        assert name not in self._vars
        self._vars[name] = var

    def __setattr__(self, name, var):
        if not name.startswith('_'):
            self._add_variable(name, var)
        else:
            # private attributes
            super(VariableHolder, self).__setattr__(name, var)

    def __getattr__(self, name):
        return self._vars[name]

    def all(self):
        """
        Returns:
            list of all variables
        """
        return list(six.itervalues(self._vars))

def _register(name, func):
    if name in _LAYER_REGISTRY:
        raise ValueError("Layer named {} is already registered!".format(name))
    if name in ['tf']:
        raise ValueError(logger.error("A layer cannot be named {}".format(name)))
    _LAYER_REGISTRY[name] = func

def get_registered_layer(name):
    """
    Args:
        name (str): the name of the layer, e.g. 'Conv2D'
    Returns:
        the wrapped layer function, or None if not registered.
    """
    return _LAYER_REGISTRY.get(name, None)

def disable_layer_logging():
    """
    Disable the shape logging for all layers from this moment on. Can be
    useful when creating multiple towers.
    """
    class ContainEverything:
        def __contains__(self, x):
            return True
    # can use nonlocal in python3, but how
    globals()['_LAYER_LOGGED'] = ContainEverything()

def layer_register(
        log_shape=False,
        use_scope=True):
    """
    Register a layer.

    Args:
        log_shape (bool): log input/output shape of this layer
        use_scope (bool or None):
            Whether to call this layer with an extra first argument as scope.
            When set to None, it can be called either with or without
            the scope name argument.
            It will try to figure out by checking if the first argument
            is string or not.
    """

    def wrapper(func):
        @wraps(func)
        def wrapped_func(*args, **kwargs):
            assert args[0] is not None, args
            if use_scope:
                name, inputs = args[0], args[1]
                args = args[1:]  # actual positional args used to call func
                assert isinstance(name, six.string_types), name
            else:
                assert not log_shape
                if isinstance(args[0], six.string_types):
                    if use_scope is False:
                        logger.warn(
                            "Please call layer {} without the first scope name argument, "
                            "or register the layer with use_scope=None to allow "
                            "two calling methods.".format(func.__name__))
                    name, inputs = args[0], args[1]
                    args = args[1:]  # actual positional args used to call func
                else:
                    inputs = args[0]
                    name = None
            if not (isinstance(inputs, (tf.Tensor, tf.Variable)) or
                    (isinstance(inputs, (list, tuple)) and
                        isinstance(inputs[0], (tf.Tensor, tf.Variable)))):
                raise ValueError("Invalid inputs to layer: " + str(inputs))

            # use kwargs from current argument scope
            actual_args = copy.copy(get_arg_scope()[func.__name__])
            # explicit kwargs overwrite argscope
            actual_args.update(kwargs)
            if six.PY3:
                # explicit positional args also override argscope. only work in PY3
                posargmap = inspect.signature(func).bind_partial(*args).arguments
                for k in six.iterkeys(posargmap):
                    if k in actual_args:
                        del actual_args[k]

            if name is not None:        # use scope
                with tf.variable_scope(name) as scope:
                    # this name is only used to surpress logging, doesn't hurt to do some heuristics
                    scope_name = re.sub('tower[0-9]+/', '', scope.name)
                    do_log_shape = log_shape and scope_name not in _LAYER_LOGGED
                    if do_log_shape:
                        logger.info("{} input: {}".format(scope.name, get_shape_str(inputs)))

                    # run the actual function
                    outputs = func(*args, **actual_args)

                    if do_log_shape:
                        # log shape info and add activation
                        logger.info("{} output: {}".format(
                            scope.name, get_shape_str(outputs)))
                        _LAYER_LOGGED.add(scope_name)
            else:
                # run the actual function
                outputs = func(*args, **actual_args)
            return outputs

        wrapped_func.symbolic_function = func   # attribute to access the underlying function object
        wrapped_func.use_scope = use_scope
        _register(func.__name__, wrapped_func)
        return wrapped_func

    return wrapper
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from collections import OrderedDict
from pprintpp import pprint as pp
import copy
import json
from Row import Row

class Dict_Sheet (object):
    def gen_alphabet_list():
        alphabet_list = []
        less26 = " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        for s in less26:
            alphabet_list.append(s)
        for i in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
            for j in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
                alphabet_list.append("%s%s" % (i, j))
        return alphabet_list
    alphabet_list = gen_alphabet_list()

    def __init__(self, wks, mapping=None):
        self.wks = wks
        """
        input:
            wks: working sheet
            kc_map: k, v mapping

        sheet1 =  {
            2: {}, <- Row
            3: {},
        }
        """

        self.__len__, self._width = self._get_len_width()
        self._id = wks._id
        if mapping is None:
            mapping = self._get_kc_map()
        self.mapping = mapping

        #self._mapping =  _mapping
        self.keys = self.mapping
        #self._update_row_mapping()

    def _get_kc_map(self):
        mapping = {}
        col_list = self.wks.row_values(1)
        for idx, value in enumerate(col_list, start=1):
            if value is not "":
                mapping[value.strip()] = idx
        return mapping

    @property
    def mapping(self):
        return self._mapping

    @mapping.setter
    def mapping(self, mapping):
        """
        input:
            wks from gspread
            mapping: {
                'col_1': 1,
                'col_2': 2,
                ...
            } or
            mapping: [
                'col_1', 'col_2', '', 'col_3'
            ]
        output: dict of key, col mapping
                ex: { id: 1, name: 2, ... }
        """
        if '_mapping' not in self.__dict__:
            _mapping = {}
        else:
            _mapping = self._mapping

        if type(mapping) is list:
            for idx, val in enumerate(mapping):
                _mapping[val] = idx

        if type(mapping) is dict:
            _mapping  = mapping

        self._mapping = _mapping
        # TO-DO upadte first rows in sheet

        last_col = len(self.wks.row_values(1))
        if last_col>26:
            last_col = 26
        #pp(last_col)
        #pp(Dict_Sheet.alphabet_list[last_col])
        first_row_range = "A1:%s1" % Dict_Sheet.alphabet_list[last_col]
        #pp(first_row_range)
        cells = self.wks.range(first_row_range)
        #cells = self.wks.range("A1:Z1")
        inv_mapping = {v: k for k, v in self._mapping.items()}
        for cell in cells:
            cell.value = ""
            if cell.col in inv_mapping:
                cell.value = inv_mapping[cell.col]
        self.wks.update_cells(cells)

        self._update_row_mapping()
        _wks_len, _wks_width = self._get_len_width()
        if _wks_width > self._width:
            self._width = _wks_width
        return _mapping

    def _get_row_cells(self, idx):
        #width = Dict_Sheet.alphabet_list[self._width]
        range_str = 'A%s:%s%s' % (str(idx), Dict_Sheet.alphabet_list[self._width], str(idx))
        #print range_str
        row_cells = self.wks.range(range_str)
        return row_cells

    def _update_row_mapping(self):
        """
        return {
            2: {},
            3: {}, ..
        }
        """
        rows = {}
        Row.kc_map[self._id] = self._mapping
        for idx in range(2, self.__len__ + 1):
            row_cells = self._get_row_cells(idx)
            row = Row(wks=self.wks, row_cells=row_cells)
            rows[idx] = row
            self.__dict__[idx] = row
        # TO-DO
        # self.__dict__[0] return k,v of row 1
        return rows

    def update(self, data_dict):
        """
        input: {
            2: { u'蜂蜜棒腿': 123, ...},
            3: { u'麻辣豬肋排': 'sss', ...}
        }
        """
        try:
            for idx, data in data_dict.items():
                row_cells = self._get_row_cells(idx)
                row = Row(wks=self.wks, row_cells=row_cells)
                row.update(data)
                self.__dict__[idx] = row
        except Exception as e:
            return -1
        return 0

    def append(self, dict_data):
        # TO-DO: append list of dict
        idx = self.__len__  + 1
        try:
            row_cells = self._get_row_cells(idx=idx)
            row = Row(wks=self.wks, row_cells=row_cells)
            row.update(dict_data)
            self.__len__ = idx
            self.__dict__[self.__len__] = row
        except Exception as e:
            raise RuntimeError('append fails')
        return row

    def __repr__(self):
        _ = copy.deepcopy(self.__dict__)
        # return json.dumps(_).encode('utf-8')
        # return repr(self.__dict__)
        return repr(_)

    def _get_len_width(self):
        all_values = self.wks.get_all_values()
        wks_len = len(all_values)
        try:
            wks_width = max([len(_) for _ in all_values])
        except Exception as e:
            wks_width = 0
        return wks_len, wks_width

    def __idx_convert__(self, idx):
        if idx >= 2:
            return idx
        elif idx < 0:
            idx = self.__len__ + idx + 1
            return idx
        else:
            raise IndexError('getitem', 'idx error')

    def __getitem__(self, key):
        if type(key) is int:
            idx = self.__idx_convert__(key)
            if idx not in self.__dict__:
                row_cells = self._get_row_cells(idx=idx)
                row = Row(wks=self.wks, row_cells=row_cells)
                self.__dict__[idx] = row
                if idx > self.__len__:
                    self.__len__ = idx
            return self.__dict__[idx]
        return self.__dict__[key]

    def __setitem__(self, key, item):
        if type(key) is int:
            idx = self.__idx_convert__(key)
            if idx == 1:
                raise IndexError('setitem', 'idx shoud not be 1')
            if idx not in self.__dict__: # Insert
                if type(item) is dict:
                    self.append(item)
                elif type(item) is Row:
                    self.__len__ += 1
                    range_str = 'A%s:%s%s' % (str(self.__len__), Dict_Sheet.alphabet_list[self._width], str(self.__len__))
                    row_list = self.wks.range(range_str)
                    row = Row(wks=self.wks, row_cells=row_list)
                    row.update(item.kv())
                    self.__dict__[self.__len__] = row
                else:
                    raise Exception('setitem', 'type error')
            else:  # Update
                row = self.__dict__[idx]
                row.clear()
                if type(item) is dict:
                    # replace
                    # ds[4] = {'LINE ID': 'test1', 'Email': 'email'}
                    #row = self.__dict__[idx]
                    row.update(item)
                    self.__dict__[idx] = row
                elif type(item) is Row:
                    #row = self.__dict__[idx]
                    row.update(item.kv())
                    self.__dict__[idx] = row
                else:
                    raise Exception('setitem', 'type error')
            if idx > self.__len__:
                self.__len__ = idx
        else:
            self.__dict__[key] = item

    def get(self, key):
        match = {}
        for row in self.rows:
            if key in row:
                for idx, item in row.items():
                    match[idx] = item
        return match

    def extend(self):
        pass

    def iteritems(self):
        for k in range(2, self.__len__+1):
            if k in self.__dict__:
                yield (k, self.__dict__[k])

    def items(self):
        for k in range(2, self.__len__+1):
            if k in self.__dict__:
                yield (k, self.__dict__[k])

    def __iter__(self):
        # TO-DO: Test __iter__ func
        # http://stackoverflow.com/questions/4019971/how-to-implement-iter-self-for-a-container-object-python
        for k in range(1, self.__len__+1):
            if k in self.__dict__:
                yield self.__dict__[k]

    def get_col_mapping(self, max_len):
        if max_len <=26:
            return " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        if max_len <=26*26:
            col_list = [' ']
            alphabets ="ABCDEFGHIJKLMNOPQRSTUVWXYZ"
            for alphabet in alphabets:
                for alphabet2 in alphabets:
                    if len(col_list) <= max_len:
                        col_list.append("%s%s" % (alphabet, alphabet2))
                    else:
                        return col_list
        raise RuntimeError('append fails')
<commit_msg>[Feature] import __future__ for print func.<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function
from collections import OrderedDict
from pprintpp import pprint as pp
import copy
import json
from Row import Row

class Dict_Sheet (object):
    def gen_alphabet_list():
        alphabet_list = []
        less26 = " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        for s in less26:
            alphabet_list.append(s)
        for i in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
            for j in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
                alphabet_list.append("%s%s" % (i, j))
        return alphabet_list
    alphabet_list = gen_alphabet_list()

    def __init__(self, wks, mapping=None):
        self.wks = wks
        """
        input:
            wks: working sheet
            kc_map: k, v mapping

        sheet1 =  {
            2: {}, <- Row
            3: {},
        }
        """

        self.__len__, self._width = self._get_len_width()
        self._id = wks._id
        if mapping is None:
            mapping = self._get_kc_map()
        self.mapping = mapping

        #self._mapping =  _mapping
        self.keys = self.mapping
        #self._update_row_mapping()

    def _get_kc_map(self):
        mapping = {}
        col_list = self.wks.row_values(1)
        for idx, value in enumerate(col_list, start=1):
            if value is not "":
                mapping[value.strip()] = idx
        return mapping

    @property
    def mapping(self):
        return self._mapping

    @mapping.setter
    def mapping(self, mapping):
        """
        input:
            wks from gspread
            mapping: {
                'col_1': 1,
                'col_2': 2,
                ...
            } or
            mapping: [
                'col_1', 'col_2', '', 'col_3'
            ]
        output: dict of key, col mapping
                ex: { id: 1, name: 2, ... }
        """
        if '_mapping' not in self.__dict__:
            _mapping = {}
        else:
            _mapping = self._mapping

        if type(mapping) is list:
            for idx, val in enumerate(mapping):
                _mapping[val] = idx

        if type(mapping) is dict:
            _mapping  = mapping

        self._mapping = _mapping
        # TO-DO upadte first rows in sheet

        last_col = len(self.wks.row_values(1))
        if last_col>26:
            last_col = 26
        #pp(last_col)
        #pp(Dict_Sheet.alphabet_list[last_col])
        first_row_range = "A1:%s1" % Dict_Sheet.alphabet_list[last_col]
        #pp(first_row_range)
        cells = self.wks.range(first_row_range)
        #cells = self.wks.range("A1:Z1")
        inv_mapping = {v: k for k, v in self._mapping.items()}
        for cell in cells:
            cell.value = ""
            if cell.col in inv_mapping:
                cell.value = inv_mapping[cell.col]
        self.wks.update_cells(cells)

        self._update_row_mapping()
        _wks_len, _wks_width = self._get_len_width()
        if _wks_width > self._width:
            self._width = _wks_width
        return _mapping

    def _get_row_cells(self, idx):
        #width = Dict_Sheet.alphabet_list[self._width]
        range_str = 'A%s:%s%s' % (str(idx), Dict_Sheet.alphabet_list[self._width], str(idx))
        #print range_str
        row_cells = self.wks.range(range_str)
        return row_cells

    def _update_row_mapping(self):
        """
        return {
            2: {},
            3: {}, ..
        }
        """
        rows = {}
        Row.kc_map[self._id] = self._mapping
        for idx in range(2, self.__len__ + 1):
            row_cells = self._get_row_cells(idx)
            row = Row(wks=self.wks, row_cells=row_cells)
            rows[idx] = row
            self.__dict__[idx] = row
        # TO-DO
        # self.__dict__[0] return k,v of row 1
        return rows

    def update(self, data_dict):
        """
        input:
            {idx: {k1: v1, k2:v2, ..}
            ex:
            {
                2: { u'蜂蜜棒腿': 123, ...},
                3: { u'麻辣豬肋排': 'sss', ...}
            }
        """
        try:
            for idx, data in data_dict.items():
                row_cells = self._get_row_cells(idx)
                row = Row(wks=self.wks, row_cells=row_cells)
                row.update(data)
                self.__dict__[idx] = row
        except Exception as e:
            return -1
        return 0

    def append(self, dict_data):
        # TO-DO: append list of dict
        idx = self.__len__  + 1
        try:
            row_cells = self._get_row_cells(idx=idx)
            row = Row(wks=self.wks, row_cells=row_cells)
            row.update(dict_data)
            self.__len__ = idx
            self.__dict__[self.__len__] = row
        except Exception as e:
            raise RuntimeError('append fails')
        return row

    def __repr__(self):
        _ = copy.deepcopy(self.__dict__)
        # return json.dumps(_).encode('utf-8')
        # return repr(self.__dict__)
        return repr(_)

    def _get_len_width(self):
        all_values = self.wks.get_all_values()
        wks_len = len(all_values)
        try:
            wks_width = max([len(_) for _ in all_values])
        except Exception as e:
            wks_width = 0
        return wks_len, wks_width

    def __idx_convert__(self, idx):
        if idx >= 2:
            return idx
        elif idx < 0:
            idx = self.__len__ + idx + 1
            return idx
        else:
            raise IndexError('getitem', 'idx error')

    def __getitem__(self, key):
        if type(key) is int:
            idx = self.__idx_convert__(key)
            if idx not in self.__dict__:
                row_cells = self._get_row_cells(idx=idx)
                row = Row(wks=self.wks, row_cells=row_cells)
                self.__dict__[idx] = row
                if idx > self.__len__:
                    self.__len__ = idx
            return self.__dict__[idx]
        return self.__dict__[key]

    def __setitem__(self, key, item):
        if type(key) is int:
            idx = self.__idx_convert__(key)
            if idx == 1:
                raise IndexError('setitem', 'idx shoud not be 1')
            if idx not in self.__dict__: # Insert
                if type(item) is dict:
                    self.append(item)
                elif type(item) is Row:
                    self.__len__ += 1
                    range_str = 'A%s:%s%s' % (str(self.__len__), Dict_Sheet.alphabet_list[self._width], str(self.__len__))
                    row_list = self.wks.range(range_str)
                    row = Row(wks=self.wks, row_cells=row_list)
                    row.update(item.kv())
                    self.__dict__[self.__len__] = row
                else:
                    raise Exception('setitem', 'type error')
            else:  # Update
                row = self.__dict__[idx]
                row.clear()
                if type(item) is dict:
                    # replace
                    # ds[4] = {'LINE ID': 'test1', 'Email': 'email'}
                    #row = self.__dict__[idx]
                    row.update(item)
                    self.__dict__[idx] = row
                elif type(item) is Row:
                    #row = self.__dict__[idx]
                    row.update(item.kv())
                    self.__dict__[idx] = row
                else:
                    raise Exception('setitem', 'type error')
            if idx > self.__len__:
                self.__len__ = idx
        else:
            self.__dict__[key] = item

    def get(self, key):
        match = {}
        for row in self.rows:
            if key in row:
                for idx, item in row.items():
                    match[idx] = item
        return match

    def extend(self):
        pass

    def iteritems(self):
        for k in range(2, self.__len__+1):
            if k in self.__dict__:
                yield (k, self.__dict__[k])

    def items(self):
        for k in range(2, self.__len__+1):
            if k in self.__dict__:
                yield (k, self.__dict__[k])

    def __iter__(self):
        # TO-DO: Test __iter__ func
        # http://stackoverflow.com/questions/4019971/how-to-implement-iter-self-for-a-container-object-python
        for k in range(1, self.__len__+1):
            if k in self.__dict__:
                yield self.__dict__[k]

    def get_col_mapping(self, max_len):
        if max_len <=26:
            return " ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        if max_len <=26*26:
            col_list = [' ']
            alphabets ="ABCDEFGHIJKLMNOPQRSTUVWXYZ"
            for alphabet in alphabets:
                for alphabet2 in alphabets:
                    if len(col_list) <= max_len:
                        col_list.append("%s%s" % (alphabet, alphabet2))
                    else:
                        return col_list
        raise RuntimeError('append fails')
<|endoftext|>


<commit_before>from __future__ import print_function
from Bio import Phylo
import imp
from pdb import set_trace
from StringIO import StringIO
import json

def modified_tree_to_json(node, extra_attr = []):
    tree_json = {}
    str_attr = ['strain','attr']
    num_attr = ['xvalue', 'yvalue', 'tvalue', 'num_date', 'clade']
    if hasattr(node, 'name'):
        tree_json['strain'] = node.name

    for prop in str_attr:
        if hasattr(node, prop):
            tree_json[prop] = node.__getattribute__(prop)
    for prop in num_attr:
        if hasattr(node, prop):
            try:
                tree_json[prop] = round(node.__getattribute__(prop),5)
            except:
                print("cannot round:", node.__getattribute__(prop), "assigned as is")
                tree_json[prop] = node.__getattribute__(prop)

    for prop in extra_attr:
        if len(prop)==2 and callable(prop[1]):
            if hasattr(node, prop[0]):
                tree_json[prop] = prop[1](node.__getattribute__(prop[0]))
        else:
            if hasattr(node, prop):
                tree_json[prop] = node.__getattribute__(prop)

    if node.clades:
        tree_json["children"] = []
        for ch in node.clades:
            tree_json["children"].append(modified_tree_to_json(ch, extra_attr))
    return tree_json

def mock_meta_json(n):
    meta = {}
    meta["updated"] = "today"
    meta["virus_count"] = n
    meta["title"] = "BEAST proof of principle (fluB)"

    meta["color_options"] = {
      "PB1": {
        "menuItem": "PB1",
        "type": "discrete",
        "legendTitle": "PB1",
        "key": "PB1"
      },
      "HA": {
        "menuItem": "HA",
        "type": "discrete",
        "legendTitle": "HA",
        "key": "HA"
      },
      "NP": {
        "menuItem": "NP",
        "type": "discrete",
        "legendTitle": "NP",
        "key": "NP"
      },
      "NA": {
        "menuItem": "NA",
        "type": "discrete",
        "legendTitle": "NA",
        "key": "NA"
      },
      "PB1.prob": {
        "menuItem": "PB1.prob",
        "type": "continuous",
        "legendTitle": "PB1.prob",
        "key": "PB1.prob"
      },
      "HA.prob": {
        "menuItem": "HA.prob",
        "type": "continuous",
        "legendTitle": "HA.prob",
        "key": "HA.prob"
      },
      "NP.prob": {
        "menuItem": "NP.prob",
        "type": "continuous",
        "legendTitle": "NP.prob",
        "key": "NP.prob"
      },
      "NA.prob": {
        "menuItem": "NA.prob",
        "type": "continuous",
        "legendTitle": "NA.prob",
        "key": "NA.prob"
      },
        "posterior": {
          "menuItem": "posterior",
          "type": "continuous",
          "legendTitle": "posterior",
          "key": "posterior"
        },
    }
    meta["filters"] = ["PB1", "HA", "NP", "NA"]
    meta["commit"] = "unknown"
    meta["panels"] = [
        "tree",
        "map",
        "entropy"
       ],
    meta["geo"] = {}
    return meta;

if __name__=="__main__":
    print("This is a proof of principle script - it should not be relied upon for analysis.")
    bt = imp.load_source('baltic', '/Users/james/blab/baltic/baltic.py')
    # set_trace()
    bt_tree=bt.loadNexus('./scratch/InfB_PB2t_ALLs1.mcc.tre',tip_regex='\_([0-9\-]+)$') ## loads a BEAST nexus file
    mostRecentTip=max([k.absoluteTime for k in bt_tree.Objects])

    simplified_tree=bt_tree.toString() ## returns a newick string without beast annotations
    xml_tree=Phylo.read(StringIO(simplified_tree),'newick') ## load simplified newick as a biopython XML tree

    xml_nodes=xml_tree.get_nonterminals(order='preorder') ## fetch nodes from XML tree
    bt_nodes=[k for k in bt_tree.traverse_tree(include_all=True) if k.branchType=='node'] ## fetch nodes from baltic tree

    xml_leaves=xml_tree.get_terminals(order='preorder') ## fetch leaves from XML tree
    bt_leaves=[k for k in bt_tree.traverse_tree()] ## fetch leaves from baltic tree

for pairs in [[xml_nodes,bt_nodes],[xml_leaves,bt_leaves]]:
    for x,b in zip(*pairs): ## iterate over pairs of branches in both trees

        setattr(x,'tvalue',b.height) ## height in years
        setattr(x,'xvalue',b.height) ## 0.0 in the future?
        setattr(x,'yvalue',b.y) ## y value after drawing
        setattr(x,'clade',b.index) ## clade

        attrs={}

        if b.absoluteTime!=None:
            attrs['num_date']=b.absoluteTime
        else:
            attrs['num_date']=bt_tree.Objects[0].absoluteTime

        if b.traits.has_key('height_95%_HPD'):
            attrs['num_date_confidence']=[mostRecentTip-t for t in b.traits['height_95%_HPD']]

        if b.branchType=='node':
            setattr(x,'name',b.index)

        for trait in b.traits:
            attrs[trait]=b.traits[trait]

        setattr(x,'attr',attrs)

    tree_json = modified_tree_to_json(xml_tree.clade)
    tree_json = tree_json["children"][0]
    json.dump(tree_json, open("scratch/beast_tree.json", 'w'), indent=2)
    meta_json = mock_meta_json(len(bt_leaves))
    json.dump(meta_json, open("scratch/beast_meta.json", 'w'), indent=2)
    json.dump({}, open("scratch/beast_sequences.json", 'w'), indent=2)

    # set_trace()

    # attrs: tree_json['children'][0]['attr'].keys()
<commit_msg>Update beast-to-auspice<commit_after>from __future__ import print_function
from Bio import Phylo
import imp
from pdb import set_trace
from StringIO import StringIO
import json

def modified_tree_to_json(node, extra_attr = []):
    tree_json = {}
    str_attr = ['strain','attr']
    num_attr = ['xvalue', 'yvalue', 'tvalue', 'num_date', 'clade']
    if hasattr(node, 'name'):
        tree_json['strain'] = node.name

    for prop in str_attr:
        if hasattr(node, prop):
            tree_json[prop] = node.__getattribute__(prop)
    for prop in num_attr:
        if hasattr(node, prop):
            try:
                tree_json[prop] = round(node.__getattribute__(prop),5)
            except:
                print("cannot round:", node.__getattribute__(prop), "assigned as is")
                tree_json[prop] = node.__getattribute__(prop)

    for prop in extra_attr:
        if len(prop)==2 and callable(prop[1]):
            if hasattr(node, prop[0]):
                tree_json[prop] = prop[1](node.__getattribute__(prop[0]))
        else:
            if hasattr(node, prop):
                tree_json[prop] = node.__getattribute__(prop)

    if node.clades:
        tree_json["children"] = []
        for ch in node.clades:
            tree_json["children"].append(modified_tree_to_json(ch, extra_attr))
    return tree_json

def mock_meta_json(n):
    meta = {}
    meta["updated"] = "today"
    meta["virus_count"] = n
    meta["title"] = "BEAST proof of principle (MERS-CoV)"

    meta["color_options"] = {
      "type": {
        "menuItem": "host",
        "type": "discrete",
        "legendTitle": "host",
        "key": "type"
      },
      "type.prob": {
        "menuItem": "host.prob",
        "type": "continuous",
        "legendTitle": "host.prob",
        "key": "type.prob"
      },
        "posterior": {
          "menuItem": "posterior",
          "type": "continuous",
          "legendTitle": "posterior",
          "key": "posterior"
        },
    }
    meta["filters"] = ["type"]
    meta["commit"] = "unknown"
    meta["panels"] = [
        "tree",
        "map",
        "entropy"
       ],
    meta["geo"] = {}
    return meta;

if __name__=="__main__":
    print("This is a proof of principle script - it should not be relied upon for analysis.")
    bt = imp.load_source('baltic', '/Users/evogytis/Documents/BLAB_baltic/baltic.py')
    # set_trace()
    bt_tree=bt.loadNexus('./scratch/MERS_274_sCoal.combinedTyped.mcc.tree',tip_regex='\|([0-9\-]+)$') ## loads a BEAST nexus file
    mostRecentTip=max([k.absoluteTime for k in bt_tree.Objects])

    simplified_tree=bt_tree.toString() ## returns a newick string without beast annotations
    xml_tree=Phylo.read(StringIO(simplified_tree),'newick') ## load simplified newick as a biopython XML tree

    xml_nodes=xml_tree.get_nonterminals(order='preorder') ## fetch nodes from XML tree
    bt_nodes=[k for k in bt_tree.traverse_tree(include_all=True) if k.branchType=='node'] ## fetch nodes from baltic tree

    xml_leaves=xml_tree.get_terminals(order='preorder') ## fetch leaves from XML tree
    bt_leaves=[k for k in bt_tree.traverse_tree()] ## fetch leaves from baltic tree

for pairs in [[xml_nodes,bt_nodes],[xml_leaves,bt_leaves]]:
    for x,b in zip(*pairs): ## iterate over pairs of branches in both trees

        setattr(x,'tvalue',b.height) ## height in years
        setattr(x,'xvalue',b.height) ## 0.0 in the future?
        setattr(x,'yvalue',b.y) ## y value after drawing
        setattr(x,'clade',b.index) ## clade

        attrs={}

        if b.absoluteTime!=None:
            attrs['num_date']=b.absoluteTime
        else:
            attrs['num_date']=bt_tree.Objects[0].absoluteTime

        if b.traits.has_key('height_95%_HPD'):
            attrs['num_date_confidence']=[mostRecentTip-t for t in b.traits['height_95%_HPD']]

        if b.branchType=='node':
            setattr(x,'name',b.index)

        for trait in b.traits:
            attrs[trait]=b.traits[trait]

        setattr(x,'attr',attrs)

    tree_json = modified_tree_to_json(xml_tree.clade)
    tree_json = tree_json["children"][0]
    json.dump(tree_json, open("scratch/beast_tree.json", 'w'), indent=2)
    meta_json = mock_meta_json(len(bt_leaves))
    json.dump(meta_json, open("scratch/beast_meta.json", 'w'), indent=2)
    json.dump({}, open("scratch/beast_sequences.json", 'w'), indent=2)

    # set_trace()

    # attrs: tree_json['children'][0]['attr'].keys()
<|endoftext|>