<commit_before>#
# ovirt-hosted-engine-setup -- ovirt hosted engine setup
# Copyright (C) 2014-2015 Red Hat, Inc.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#

"""
NFS / GlusterFS storage plugin.
"""

import gettext
import os
import tempfile
import time
import xml.dom.minidom

from otopi import plugin
from otopi import util

from ovirt_hosted_engine_setup import constants as ohostedcons
from ovirt_hosted_engine_setup import domains as ohosteddomains

def _(m):
    return gettext.dgettext(message=m, domain='ovirt-hosted-engine-setup')

@util.export
class Plugin(plugin.PluginBase):
    """
    NFS / GlusterFS storage plugin.
    """

    UMOUNT_TRIES = 10

    def __init__(self, context):
        super(Plugin, self).__init__(context=context)
        self._checker = ohosteddomains.DomainChecker()

    def _mount(self, path, connection, domain_type):
        fstype = ''
        opts = []

        if domain_type == ohostedcons.DomainTypes.NFS3:
            fstype = ohostedcons.FileSystemTypes.NFS
            opts.append('vers=3')
        elif domain_type == ohostedcons.DomainTypes.NFS4:
            fstype = ohostedcons.FileSystemTypes.NFS
            opts.append('vers=4')
        elif domain_type == ohostedcons.DomainTypes.GLUSTERFS:
            fstype = ohostedcons.FileSystemTypes.GLUSTERFS

        if fstype == ohostedcons.FileSystemTypes.NFS:
            opts.append('retry=1')

        mount_cmd = (
            self.command.get('mount'),
            '-t%s' % fstype,
        )

        if opts:
            mount_cmd += (
                '-o%s' % ','.join(opts),
            )

        mount_cmd += (
            connection,
            path,
        )

        rc, _stdout, stderr = self.execute(
            mount_cmd,
            raiseOnError=False,
            env={
                'LC_ALL': 'C',
            },

        )
        error = '\n'.join(stderr)
        if rc != 0:
            self.logger.error(
                _(
                    'Error while mounting specified storage path: {error}'
                ).format(
                    error=error,
                )
            )
            raise RuntimeError(error)

    def _umount(self, path):
        rc = -1
        tries = self.UMOUNT_TRIES
        while tries > 0:
            rc, _stdout, _stderr = self.execute(
                (
                    self.command.get('umount'),
                    path
                ),
                raiseOnError=False,
                env={
                    'LC_ALL': 'C',
                },
            )
            if rc == 0:
                tries = -1
            else:
                tries -= 1
                time.sleep(1)
                # rc, stdout and stderr are automatically logged as debug
                self.execute(
                    (
                        self.command.get('lsof'),
                        '+D%s' % path,
                        '-xfl'
                    ),
                    raiseOnError=False,
                    env={
                        'LC_ALL': 'C',
                    },
                )
        return rc

    def _check_domain_rights(self, path):
        rc, _stdout, _stderr = self.execute(
            (
                self.command.get('sudo'),
                '-u', 'vdsm',
                '-g', 'kvm',
                'test',
                '-r', path,
                '-a',
                '-w', path,
                '-a',
                '-x', path,
            ),
            raiseOnError=False
        )
        if rc != 0:
            raise RuntimeError(
                _(
                    'permission settings on the specified storage do not '
                    'allow access to the storage to vdsm user and kvm group. '
                    'Verify permission settings on the specified storage '
                    'or specify another location'
                )
            )

    def _check_volume_properties(self, connection):
        server, volume = connection.split(':')
        if volume[0] == '/':
            volume = volume[1:]
        glustercmd = self.command.get('gluster')
        rc, stdout, stderr = self.execute(
            args=(
                glustercmd,
                '--mode=script',
                '--xml',
                'volume',
                'info',
                volume,
                '--remote-host=%s' % server,
            ),
            raiseOnError=True
        )
        if rc != 0:
            self.logger.error(_('Failed to retrieve Gluster Volume info'))
            raise RuntimeError(
                'Failed to retrieve Gluster Volume info: ' + str(stdout)
            )
        self.logger.debug('Gluster Volume info XML output: ' + str(stdout))
        dom = xml.dom.minidom.parseString(''.join(stdout))
        cliOutput = dom.getElementsByTagName('cliOutput')[0]
        opRet = (
            cliOutput.getElementsByTagName('opRet')[0]
        ).firstChild.nodeValue
        opErrno = (
            cliOutput.getElementsByTagName('opErrno')[0]
        ).firstChild.nodeValue
        opErrstrElement = cliOutput.getElementsByTagName('opErrstr')[0]
        opErrstr = ' '.join(
            t.nodeValue for t in opErrstrElement.childNodes
            if t.nodeType == t.TEXT_NODE
        )

        if opRet != '0':
            self.logger.error(_('Failed to retrieve Gluster Volume info'))
            raise RuntimeError(
                'Failed to retrieve Gluster Volume info '
                '[{code}]: {error}'.format(
                    code=opErrno,
                    error=opErrstr,
                )
            )
        volInfo = cliOutput.getElementsByTagName('volInfo')[0]
        volumes = volInfo.getElementsByTagName('volumes')[0]
        volCount = (
            volumes.getElementsByTagName('count')[0]
        ).firstChild.nodeValue
        if volCount != '1':
            raise RuntimeError(
                _('GlusterFS Volume {volume} does not exist!').format(
                    volume=volume,
                )
            )
        volume = volumes.getElementsByTagName('volume')[0]
        replicaCount = (
            volume.getElementsByTagName('replicaCount')[0]
        ).firstChild.nodeValue
        if replicaCount != '3':
            raise RuntimeError(
                _(
                    'GlusterFS Volume is not using replica 3'
                )
            )
        self.logger.info(_('GlusterFS replica 3 Volume detected'))
        host_list = []
        for brickE in volume.getElementsByTagName('brick'):
            host_list.append(
                (brickE.getElementsByTagName('hostUuid')[0])
                .firstChild.nodeValue
            )
        if len(set(host_list)) < 3:
            self.logger.warning(_(
                'Three distinct hosts are required for '
                'safe and reliable operations'
            ))
        status = (
            volume.getElementsByTagName('status')[0]
        ).firstChild.nodeValue
        statusStr = (
            volume.getElementsByTagName('statusStr')[0]
        ).firstChild.nodeValue
        if status != '1':
            raise RuntimeError(
                _(
                    "GlusterFS Volume is '{statusStr}', "
                    "please ensure that it's started"
                ).format(
                    statusStr=statusStr,
                )
            )

    def _validateDomain(self, connection, domain_type, check_space):
        if self.environment[
            ohostedcons.StorageEnv.DOMAIN_TYPE
        ] == ohostedcons.DomainTypes.GLUSTERFS:
            self._check_volume_properties(connection)
        path = tempfile.mkdtemp()
        try:
            self._mount(path, connection, domain_type)
            self._checker.check_valid_path(path)
            self._check_domain_rights(path)
            self._checker.check_base_writable(path)
            if check_space:
                self._checker.check_available_space(
                    path,
                    ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB
                )
        finally:
            if self._umount(path) == 0:
                os.rmdir(path)
            else:
                self.logger.warning(
                    _('Cannot unmount {path}').format(
                        path=path,
                    )
                )

    @plugin.event(
        stage=plugin.Stages.STAGE_SETUP,
    )
    def _setup(self):
        self.command.detect('lsof')
        self.command.detect('sudo')
        self.command.detect('mount')
        self.command.detect('umount')
        self.command.detect('gluster')

    @plugin.event(
        stage=plugin.Stages.STAGE_CUSTOMIZATION,
        name=ohostedcons.Stages.CONFIG_STORAGE_NFS,
        after=(
            ohostedcons.Stages.CONFIG_STORAGE_EARLY,
        ),
        before=(
            ohostedcons.Stages.CONFIG_STORAGE_LATE,
        ),
        condition=lambda self: (
            self.environment[ohostedcons.StorageEnv.DOMAIN_TYPE] in (
                ohostedcons.DomainTypes.GLUSTERFS,
                ohostedcons.DomainTypes.NFS3,
                ohostedcons.DomainTypes.NFS4,
            ) and
            not self.environment[
                ohostedcons.StorageEnv.GLUSTER_PROVISIONING_ENABLED
            ]
        ),
    )
    def _customization(self):
        if self.environment[
            ohostedcons.StorageEnv.DOMAIN_TYPE
        ] == ohostedcons.DomainTypes.GLUSTERFS:
            self.logger.info(
                _(
                    'Please note that Replica 3 support is required for '
                    'the shared storage.'
                )
            )
        interactive = self.environment[
            ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
        ] is None
        validDomain = False
        while not validDomain:
            if interactive:
                self.environment[
                    ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                ] = self.dialog.queryString(
                    name='OVEHOSTED_STORAGE_DOMAIN_CONNECTION',
                    note=_(
                        'Please specify the full shared storage '
                        'connection path to use (example: host:/path): '
                    ),
                    prompt=True,
                    caseSensitive=True,
                )
            try:
                self._validateDomain(
                    connection=self.environment[
                        ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                    ],
                    domain_type=self.environment[
                        ohostedcons.StorageEnv.DOMAIN_TYPE
                    ],
                    check_space=False,
                )
                validDomain = True
            except (ValueError, RuntimeError) as e:
                if interactive:
                    self.logger.debug('exception', exc_info=True)
                    self.logger.error(
                        _(
                            'Cannot access storage connection '
                            '{connection}: {error}'
                        ).format(
                            connection=self.environment[
                                ohostedcons.StorageEnv.
                                STORAGE_DOMAIN_CONNECTION
                            ],
                            error=e,
                        )
                    )
                else:
                    raise e
            except ohosteddomains.InsufficientSpaceError as e:
                self.logger.debug('exception', exc_info=True)
                self.logger.debug(e)
                min_requirement = '%0.2f' % (
                    ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB / 1024.0
                )
                if interactive:
                    self.logger.error(
                        _(
                            'Storage domain for self hosted engine '
                            'is too small: '
                            'you should have at least {min_r} GB free'.format(
                                min_r=min_requirement,
                            )
                        )
                    )
                else:
                    raise RuntimeError(
                        _(
                            'Storage domain for self hosted engine '
                            'is too small: '
                            'you should have at least {min_r} GB free'.format(
                                min_r=min_requirement,
                            )
                        )
                    )

    @plugin.event(
        stage=plugin.Stages.STAGE_CUSTOMIZATION,
        after=(
            ohostedcons.Stages.CONFIG_STORAGE_LATE,
        ),
        before=(
            ohostedcons.Stages.DIALOG_TITLES_E_STORAGE,
        ),
        condition=lambda self: (
            not self.environment[
                ohostedcons.CoreEnv.IS_ADDITIONAL_HOST
            ] and
            self.environment[ohostedcons.StorageEnv.DOMAIN_TYPE] in (
                ohostedcons.DomainTypes.GLUSTERFS,
                ohostedcons.DomainTypes.NFS3,
                ohostedcons.DomainTypes.NFS4,
            ) and
            not self.environment[
                ohostedcons.StorageEnv.GLUSTER_PROVISIONING_ENABLED
            ]
        ),
    )
    def _late_customization(self):
        # On first host we need to check if we have enough space too.
        # We must skip this check on additional hosts because the space is
        # already filled with the Hosted Engine VM image.
        # Sadly we can't go back to previous customization stage so here
        # we can only fail the setup.
        try:
            self._validateDomain(
                connection=self.environment[
                    ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                ],
                domain_type=self.environment[
                    ohostedcons.StorageEnv.DOMAIN_TYPE
                ],
                check_space=True,
            )
        except ohosteddomains.InsufficientSpaceError as e:
            self.logger.debug('exception', exc_info=True)
            self.logger.debug(e)
            min_requirement = '%0.2f' % (
                ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB / 1024.0
            )
            raise RuntimeError(
                _(
                    'Storage domain for self hosted engine '
                    'is too small: '
                    'you should have at least {min_r} GB free'.format(
                        min_r=min_requirement,
                    )
                )
            )

# vim: expandtab tabstop=4 shiftwidth=4
<commit_msg>packaging: setup: validating export paths as VDSM does<commit_after>#
# ovirt-hosted-engine-setup -- ovirt hosted engine setup
# Copyright (C) 2014-2015 Red Hat, Inc.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#

"""
NFS / GlusterFS storage plugin.
"""

import gettext
import os
import tempfile
import time
import xml.dom.minidom

from otopi import plugin
from otopi import util

from ovirt_hosted_engine_setup import constants as ohostedcons
from ovirt_hosted_engine_setup import domains as ohosteddomains

def _(m):
    return gettext.dgettext(message=m, domain='ovirt-hosted-engine-setup')

@util.export
class Plugin(plugin.PluginBase):
    """
    NFS / GlusterFS storage plugin.
    """

    UMOUNT_TRIES = 10

    def __init__(self, context):
        super(Plugin, self).__init__(context=context)
        self._checker = ohosteddomains.DomainChecker()

    def _mount(self, path, connection, domain_type):
        fstype = ''
        opts = []

        if domain_type == ohostedcons.DomainTypes.NFS3:
            fstype = ohostedcons.FileSystemTypes.NFS
            opts.append('vers=3')
        elif domain_type == ohostedcons.DomainTypes.NFS4:
            fstype = ohostedcons.FileSystemTypes.NFS
            opts.append('vers=4')
        elif domain_type == ohostedcons.DomainTypes.GLUSTERFS:
            fstype = ohostedcons.FileSystemTypes.GLUSTERFS

        if fstype == ohostedcons.FileSystemTypes.NFS:
            opts.append('retry=1')

        mount_cmd = (
            self.command.get('mount'),
            '-t%s' % fstype,
        )

        if opts:
            mount_cmd += (
                '-o%s' % ','.join(opts),
            )

        mount_cmd += (
            connection,
            path,
        )

        rc, _stdout, stderr = self.execute(
            mount_cmd,
            raiseOnError=False,
            env={
                'LC_ALL': 'C',
            },

        )
        error = '\n'.join(stderr)
        if rc != 0:
            self.logger.error(
                _(
                    'Error while mounting specified storage path: {error}'
                ).format(
                    error=error,
                )
            )
            raise RuntimeError(error)

    def _umount(self, path):
        rc = -1
        tries = self.UMOUNT_TRIES
        while tries > 0:
            rc, _stdout, _stderr = self.execute(
                (
                    self.command.get('umount'),
                    path
                ),
                raiseOnError=False,
                env={
                    'LC_ALL': 'C',
                },
            )
            if rc == 0:
                tries = -1
            else:
                tries -= 1
                time.sleep(1)
                # rc, stdout and stderr are automatically logged as debug
                self.execute(
                    (
                        self.command.get('lsof'),
                        '+D%s' % path,
                        '-xfl'
                    ),
                    raiseOnError=False,
                    env={
                        'LC_ALL': 'C',
                    },
                )
        return rc

    def _check_domain_rights(self, path):
        rc, _stdout, _stderr = self.execute(
            (
                self.command.get('sudo'),
                '-u', 'vdsm',
                '-g', 'kvm',
                'test',
                '-r', path,
                '-a',
                '-w', path,
                '-a',
                '-x', path,
            ),
            raiseOnError=False
        )
        if rc != 0:
            raise RuntimeError(
                _(
                    'permission settings on the specified storage do not '
                    'allow access to the storage to vdsm user and kvm group. '
                    'Verify permission settings on the specified storage '
                    'or specify another location'
                )
            )

    def _check_volume_properties(self, connection):
        server, volume = connection.split(':')
        if volume[0] == '/':
            volume = volume[1:]
        glustercmd = self.command.get('gluster')
        rc, stdout, stderr = self.execute(
            args=(
                glustercmd,
                '--mode=script',
                '--xml',
                'volume',
                'info',
                volume,
                '--remote-host=%s' % server,
            ),
            raiseOnError=True
        )
        if rc != 0:
            self.logger.error(_('Failed to retrieve Gluster Volume info'))
            raise RuntimeError(
                'Failed to retrieve Gluster Volume info: ' + str(stdout)
            )
        self.logger.debug('Gluster Volume info XML output: ' + str(stdout))
        dom = xml.dom.minidom.parseString(''.join(stdout))
        cliOutput = dom.getElementsByTagName('cliOutput')[0]
        opRet = (
            cliOutput.getElementsByTagName('opRet')[0]
        ).firstChild.nodeValue
        opErrno = (
            cliOutput.getElementsByTagName('opErrno')[0]
        ).firstChild.nodeValue
        opErrstrElement = cliOutput.getElementsByTagName('opErrstr')[0]
        opErrstr = ' '.join(
            t.nodeValue for t in opErrstrElement.childNodes
            if t.nodeType == t.TEXT_NODE
        )

        if opRet != '0':
            self.logger.error(_('Failed to retrieve Gluster Volume info'))
            raise RuntimeError(
                'Failed to retrieve Gluster Volume info '
                '[{code}]: {error}'.format(
                    code=opErrno,
                    error=opErrstr,
                )
            )
        volInfo = cliOutput.getElementsByTagName('volInfo')[0]
        volumes = volInfo.getElementsByTagName('volumes')[0]
        volCount = (
            volumes.getElementsByTagName('count')[0]
        ).firstChild.nodeValue
        if volCount != '1':
            raise RuntimeError(
                _('GlusterFS Volume {volume} does not exist!').format(
                    volume=volume,
                )
            )
        volume = volumes.getElementsByTagName('volume')[0]
        replicaCount = (
            volume.getElementsByTagName('replicaCount')[0]
        ).firstChild.nodeValue
        if replicaCount != '3':
            raise RuntimeError(
                _(
                    'GlusterFS Volume is not using replica 3'
                )
            )
        self.logger.info(_('GlusterFS replica 3 Volume detected'))
        host_list = []
        for brickE in volume.getElementsByTagName('brick'):
            host_list.append(
                (brickE.getElementsByTagName('hostUuid')[0])
                .firstChild.nodeValue
            )
        if len(set(host_list)) < 3:
            self.logger.warning(_(
                'Three distinct hosts are required for '
                'safe and reliable operations'
            ))
        status = (
            volume.getElementsByTagName('status')[0]
        ).firstChild.nodeValue
        statusStr = (
            volume.getElementsByTagName('statusStr')[0]
        ).firstChild.nodeValue
        if status != '1':
            raise RuntimeError(
                _(
                    "GlusterFS Volume is '{statusStr}', "
                    "please ensure that it's started"
                ).format(
                    statusStr=statusStr,
                )
            )

    def _fix_path_syntax(self):
        path = self.environment[
            ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
        ]
        if (
            path and
            len(path) > 2 and
            path[-1] == '/' and
            path[-2] == ':'
        ):
            raise RuntimeError(
                _(
                    "'server:/' is not an acceptable export path, "
                    "please use at least one subdirectory"
                )
            )
        path = os.path.normpath(
            self.environment[
                ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
            ]
        )
        if path != self.environment[
            ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
        ]:
            self.logger.warning(
                _(
                    "Fixing path syntax: "
                    "replacing '{original}' with '{fixed}'"
                ).format(
                    original=self.environment[
                        ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                    ],
                    fixed=path,
                )
            )
            self.environment[
                ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
            ] = path

    def _validateDomain(self, connection, domain_type, check_space):
        if self.environment[
            ohostedcons.StorageEnv.DOMAIN_TYPE
        ] == ohostedcons.DomainTypes.GLUSTERFS:
            self._check_volume_properties(connection)
        path = tempfile.mkdtemp()
        try:
            self._mount(path, connection, domain_type)
            self._checker.check_valid_path(path)
            self._check_domain_rights(path)
            self._checker.check_base_writable(path)
            if check_space:
                self._checker.check_available_space(
                    path,
                    ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB
                )
        finally:
            if self._umount(path) == 0:
                os.rmdir(path)
            else:
                self.logger.warning(
                    _('Cannot unmount {path}').format(
                        path=path,
                    )
                )

    @plugin.event(
        stage=plugin.Stages.STAGE_SETUP,
    )
    def _setup(self):
        self.command.detect('lsof')
        self.command.detect('sudo')
        self.command.detect('mount')
        self.command.detect('umount')
        self.command.detect('gluster')

    @plugin.event(
        stage=plugin.Stages.STAGE_CUSTOMIZATION,
        name=ohostedcons.Stages.CONFIG_STORAGE_NFS,
        after=(
            ohostedcons.Stages.CONFIG_STORAGE_EARLY,
        ),
        before=(
            ohostedcons.Stages.CONFIG_STORAGE_LATE,
        ),
        condition=lambda self: (
            self.environment[ohostedcons.StorageEnv.DOMAIN_TYPE] in (
                ohostedcons.DomainTypes.GLUSTERFS,
                ohostedcons.DomainTypes.NFS3,
                ohostedcons.DomainTypes.NFS4,
            ) and
            not self.environment[
                ohostedcons.StorageEnv.GLUSTER_PROVISIONING_ENABLED
            ]
        ),
    )
    def _customization(self):
        if self.environment[
            ohostedcons.StorageEnv.DOMAIN_TYPE
        ] == ohostedcons.DomainTypes.GLUSTERFS:
            self.logger.info(
                _(
                    'Please note that Replica 3 support is required for '
                    'the shared storage.'
                )
            )
        interactive = self.environment[
            ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
        ] is None
        validDomain = False
        while not validDomain:
            if interactive:
                self.environment[
                    ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                ] = self.dialog.queryString(
                    name='OVEHOSTED_STORAGE_DOMAIN_CONNECTION',
                    note=_(
                        'Please specify the full shared storage '
                        'connection path to use (example: host:/path): '
                    ),
                    prompt=True,
                    caseSensitive=True,
                )
            try:
                self._fix_path_syntax()
                self._validateDomain(
                    connection=self.environment[
                        ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                    ],
                    domain_type=self.environment[
                        ohostedcons.StorageEnv.DOMAIN_TYPE
                    ],
                    check_space=False,
                )
                validDomain = True
            except (ValueError, RuntimeError) as e:
                if interactive:
                    self.logger.debug('exception', exc_info=True)
                    self.logger.error(
                        _(
                            'Cannot access storage connection '
                            '{connection}: {error}'
                        ).format(
                            connection=self.environment[
                                ohostedcons.StorageEnv.
                                STORAGE_DOMAIN_CONNECTION
                            ],
                            error=e,
                        )
                    )
                else:
                    raise e
            except ohosteddomains.InsufficientSpaceError as e:
                self.logger.debug('exception', exc_info=True)
                self.logger.debug(e)
                min_requirement = '%0.2f' % (
                    ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB / 1024.0
                )
                if interactive:
                    self.logger.error(
                        _(
                            'Storage domain for self hosted engine '
                            'is too small: '
                            'you should have at least {min_r} GB free'.format(
                                min_r=min_requirement,
                            )
                        )
                    )
                else:
                    raise RuntimeError(
                        _(
                            'Storage domain for self hosted engine '
                            'is too small: '
                            'you should have at least {min_r} GB free'.format(
                                min_r=min_requirement,
                            )
                        )
                    )

    @plugin.event(
        stage=plugin.Stages.STAGE_CUSTOMIZATION,
        after=(
            ohostedcons.Stages.CONFIG_STORAGE_LATE,
        ),
        before=(
            ohostedcons.Stages.DIALOG_TITLES_E_STORAGE,
        ),
        condition=lambda self: (
            not self.environment[
                ohostedcons.CoreEnv.IS_ADDITIONAL_HOST
            ] and
            self.environment[ohostedcons.StorageEnv.DOMAIN_TYPE] in (
                ohostedcons.DomainTypes.GLUSTERFS,
                ohostedcons.DomainTypes.NFS3,
                ohostedcons.DomainTypes.NFS4,
            ) and
            not self.environment[
                ohostedcons.StorageEnv.GLUSTER_PROVISIONING_ENABLED
            ]
        ),
    )
    def _late_customization(self):
        # On first host we need to check if we have enough space too.
        # We must skip this check on additional hosts because the space is
        # already filled with the Hosted Engine VM image.
        # Sadly we can't go back to previous customization stage so here
        # we can only fail the setup.
        try:
            self._validateDomain(
                connection=self.environment[
                    ohostedcons.StorageEnv.STORAGE_DOMAIN_CONNECTION
                ],
                domain_type=self.environment[
                    ohostedcons.StorageEnv.DOMAIN_TYPE
                ],
                check_space=True,
            )
        except ohosteddomains.InsufficientSpaceError as e:
            self.logger.debug('exception', exc_info=True)
            self.logger.debug(e)
            min_requirement = '%0.2f' % (
                ohostedcons.Const.MINIMUM_SPACE_STORAGEDOMAIN_MB / 1024.0
            )
            raise RuntimeError(
                _(
                    'Storage domain for self hosted engine '
                    'is too small: '
                    'you should have at least {min_r} GB free'.format(
                        min_r=min_requirement,
                    )
                )
            )

# vim: expandtab tabstop=4 shiftwidth=4
<|endoftext|>


<commit_before># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from numbers import Number

import statsd
from oslo_config import cfg

from st2common.metrics.base import BaseMetricsDriver
from st2common.metrics.utils import check_key
from st2common.metrics.utils import get_full_key_name

__all__ = [
    'StatsdDriver'
]

class StatsdDriver(BaseMetricsDriver):
    """
    StatsD Implementation of the metrics driver
    """
    def __init__(self):
        statsd.Connection.set_defaults(host=cfg.CONF.metrics.host, port=cfg.CONF.metrics.port,
                                       sample_rate=cfg.CONF.metrics.sample_rate)

        self._counters = {}
        self._timer = statsd.Timer('')

    def time(self, key, time):
        """
        Timer metric
        """
        check_key(key)
        assert isinstance(time, Number)

        key = get_full_key_name(key)
        self._timer.send(key, time)

    def inc_counter(self, key, amount=1):
        """
        Increment counter
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        self._counters[key] = self._counters.get(key, statsd.Counter(key))
        self._counters[key] += amount

    def dec_counter(self, key, amount=1):
        """
        Decrement metric
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        self._counters[key] = self._counters.get(key, statsd.Counter(key))
        self._counters[key] -= amount

    def set_gauge(self, key, value):
        """
        Set gauge value.
        """
        check_key(key)
        assert isinstance(value, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.send(None, value)

    def inc_gauge(self, key, amount=1):
        """
        Increment gauge value.
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.increment(None, amount)

    def dec_gauge(self, key, amount=1):
        """
        Decrement gauge value.
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.decrement(None, amount)
<commit_msg>Instead of caching Counter objects in memory, always create a new one.<commit_after># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from numbers import Number

import statsd
from oslo_config import cfg

from st2common.metrics.base import BaseMetricsDriver
from st2common.metrics.utils import check_key
from st2common.metrics.utils import get_full_key_name

__all__ = [
    'StatsdDriver'
]

class StatsdDriver(BaseMetricsDriver):
    """
    StatsD Implementation of the metrics driver
    """
    def __init__(self):
        statsd.Connection.set_defaults(host=cfg.CONF.metrics.host, port=cfg.CONF.metrics.port,
                                       sample_rate=cfg.CONF.metrics.sample_rate)

    def time(self, key, time):
        """
        Timer metric
        """
        check_key(key)
        assert isinstance(time, Number)

        key = get_full_key_name(key)
        timer = statsd.Timer('')
        timer.send(key, time)

    def inc_counter(self, key, amount=1):
        """
        Increment counter
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)

        counter = statsd.Counter(key)
        counter.increment(delta=amount)

    def dec_counter(self, key, amount=1):
        """
        Decrement metric
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        counter = statsd.Counter(key)
        counter.decrement(delta=amount)

    def set_gauge(self, key, value):
        """
        Set gauge value.
        """
        check_key(key)
        assert isinstance(value, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.send(None, value)

    def inc_gauge(self, key, amount=1):
        """
        Increment gauge value.
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.increment(None, amount)

    def dec_gauge(self, key, amount=1):
        """
        Decrement gauge value.
        """
        check_key(key)
        assert isinstance(amount, Number)

        key = get_full_key_name(key)
        gauge = statsd.Gauge(key)
        gauge.decrement(None, amount)
<|endoftext|>


<commit_before>from tendrl.commons.event import Event
from tendrl.commons.message import ExceptionMessage
from tendrl.commons.utils import log_utils as logger

def sync(sync_ttl):
    try:
        if NS.tendrl_context.integration_id:
            if "provisioner/%s" % NS.tendrl_context.integration_id \
                in NS.node_context.tags:
                logger.log(
                    "debug",
                    NS.publisher_id,
                    {"message": "node_sync, updating "
                                "cluster tendrl context"}
                )
                NS.tendrl.objects.ClusterTendrlContext(
                    integration_id=NS.tendrl_context.integration_id,
                    cluster_id=NS.tendrl_context.cluster_id,
                    cluster_name=NS.tendrl_context.cluster_name,
                    sds_name=NS.tendrl_context.sds_name,
                    sds_version=NS.tendrl_context.sds_version
                ).save()

            logger.log(
                "debug",
                NS.publisher_id,
                {"message": "node_sync, Updating"
                            "cluster node context"}
            )
            NS.tendrl.objects.ClusterNodeContext(
                node_id=NS.node_context.node_id,
                fqdn=NS.node_context.fqdn,
                status="UP",
                tags=NS.node_context.tags
            ).save(ttl=sync_ttl)
    except Exception as ex:
        Event(
            ExceptionMessage(
                priority="error",
                publisher=NS.publisher_id,
                payload={"message": "node_sync cluster contexts sync "
                                    "failed: " + ex.message,
                         "exception": ex}
            )
        )
<commit_msg>Update cluster_contexts_sync.py<commit_after>from tendrl.commons.event import Event
from tendrl.commons.message import ExceptionMessage
from tendrl.commons.utils import log_utils as logger

def sync(sync_ttl):
    try:
        if NS.tendrl_context.integration_id:
            if "provisioner/%s" % NS.tendrl_context.integration_id \
                in NS.node_context.tags:
                logger.log(
                    "debug",
                    NS.publisher_id,
                    {"message": "node_sync, updating "
                                "cluster tendrl context"}
                )
                NS.tendrl.objects.ClusterTendrlContext(
                    integration_id=NS.tendrl_context.integration_id,
                    cluster_id=NS.tendrl_context.cluster_id,
                    cluster_name=NS.tendrl_context.cluster_name,
                    sds_name=NS.tendrl_context.sds_name,
                    sds_version=NS.tendrl_context.sds_version
                ).save()

            logger.log(
                "debug",
                NS.publisher_id,
                {"message": "node_sync, Updating"
                            "cluster node context"}
            )
            NS.tendrl.objects.ClusterNodeContext(
                node_id=NS.node_context.node_id,
                fqdn=NS.node_context.fqdn,
                status="UP",
                tags=NS.node_context.tags
            ).save(ttl=sync_ttl)
    except Exception as ex:
        Event(
            ExceptionMessage(
                priority="error",
                publisher=NS.publisher_id,
                payload={"message": "node_sync cluster contexts sync "
                                    "failed: " + ex.message,
                         "exception": ex}
            )
        )
<|endoftext|>


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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.contrib import tensorrt as trt
from tensorflow.core.protobuf import config_pb2 as cpb2
from tensorflow.python.client import session as csess
from tensorflow.python.framework import test_util
from tensorflow.python.framework import constant_op as cop
from tensorflow.python.framework import dtypes as dtypes
from tensorflow.python.framework import importer as importer
from tensorflow.python.framework import ops as ops
from tensorflow.python.ops import array_ops as aops
from tensorflow.python.ops import nn as nn
from tensorflow.python.ops import nn_ops as nn_ops
from tensorflow.python.platform import googletest
from tensorflow.python.platform import test

@test_util.with_c_api
class IntegrationTest(test_util.TensofFlowTestCase):

  def setUp(self):
    """ Setup method """
    super(IntegrationTest, self).setUp()
    warnings.simplefilter('always')
    inp_dims = (100, 24, 24, 2)
    self._input = np.random.random_sample(inp_dims)
    self._original_graph = get_simple_graph_def()
    self._gpu_options = cpb2.GPUOptions(
        per_process_gpu_memory_fraction=0.50)
    self._config = cpb2.ConfigProto(gpu_options=gpu_options)
    self._reference = self.run_graph(self._original_graph, self._input)

  def get_simple_graph_def(self):
    """Create a simple graph and return its graph_def."""
    g = ops.Graph()
    with g.as_default():
      a = aops.placeholder(
          dtype=dtypes.float32, shape=(None, 24, 24, 2), name="input")
      e = cop.constant(
          [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
          name="weights",
          dtype=dtypes.float32)
      conv = nn.conv2d(
          input=a,
          filter=e,
          strides=[1, 2, 2, 1],
          padding="SAME",
          name="conv")
      b = cop.constant(
          [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtypes.float32)
      t = nn.bias_add(conv, b, name="biasAdd")
      relu = nn.relu(t, "relu")
      idty = aops.identity(relu, "ID")
      v = nn_ops.max_pool(
          idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
      aops.squeeze(v, name="output")
    return g.as_graph_def()

  def run_graph(self, gdef, dumm_inp):
    """Run given graphdef once."""
    ops.reset_default_graph()
    g = ops.Graph()
    with g.as_default():
      inp, out = importer.import_graph_def(
          graph_def=gdef, return_elements=["input", "output"])
      inp = inp.outputs[0]
      out = out.outputs[0]
    with self.test_session(
        grap=g, config=self._config, use_gpu=True,
        force_gpu=True) as sess:
      val = sess.run(out, {inp: dumm_inp})
    return val

  # Use real data that is representative of the inference dataset
  # for calibration. For this test script it is random data.
  def run_calibration(self, gdef, dumm_inp):
    """Run given calibration graph multiple times."""
    ops.reset_default_graph()
    g = ops.Graph()
    with g.as_default():
      inp, out = importer.import_graph_def(
          graph_def=gdef, return_elements=["input", "output"])
      inp = inp.outputs[0]
      out = out.outputs[0]
      # run over real calibration data here, we are mimicking a calibration set of
      # 30 different batches. Use as much calibration data as you want
    with self.test_session(
        grap=g, config=self._config, use_gpu=True,
        force_gpu=True) as sess:
      for _ in range(30):
        val = sess.run(out, {inp: dumm_inp})
    return val

  def get_trt_graph(self, mode):
    """  return trt converted graph """
    if mode == "FP32":
      return trt.create_inference_graph(
          input_graph_def=self._orig_graph,
          outputs=["output"],
          max_batch_size=inp_dims[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "FP32",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )
    elif mode == "FP16":
      return trt.create_inference_graph(
          input_graph_def=self._orig_graph,
          outputs=["output"],
          max_batch_size=inp_dims[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "FP16",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )
    elif mode == "INT8":
      return trt.create_inference_graph(
          input_graph_def=self._orig_graph,
          outputs=["output"],
          max_batch_size=inp_dims[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "INT8",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )

    return None

  def testFP32(self):
    """ Test FP32 conversion. Results should be identical to native case """
    trt_graph = self.get_trt_graph("FP32")
    result = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(self._reference, result)
    result = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(self._reference, result)

  def testFP16(self):
    """ Test FP16 conversion. Results may be different from native case """
    trt_graph = self.get_trt_graph("FP16")
    result = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(self._reference, result)
    result = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(self._reference, result)

  def testINT8(self):
    """ Test INT8 conversion. Results may be different from native case """
    calib_graph = self.get_trt_graph("INT8")
    result = self.run_calibration(calib_graph, self._input)
    self.assertAllEqual(self._reference, result)
    int8_graph = trt.calib_graph_to_infer_graph(int8_calib_gdef)
    result = self.run_graph(int8_graph, self._input)
    self.assertAllEqual(self._reference, result)
    result = self.run_graph(int8_graph, self._input)
    self.assertAllEqual(self._reference, result)

if __name__ == '__main__':
  googletest.main()
<commit_msg>Fix testing<commit_after># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Script to test TF-TensorRT integration."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import warnings

from tensorflow.contrib import tensorrt as trt
from tensorflow.core.protobuf import config_pb2 as cpb2
from tensorflow.python.client import session as csess
from tensorflow.python.framework import test_util
from tensorflow.python.framework import constant_op as cop
from tensorflow.python.framework import dtypes as dtypes
from tensorflow.python.framework import importer as importer
from tensorflow.python.framework import ops as ops
from tensorflow.python.ops import array_ops as aops
from tensorflow.python.ops import nn as nn
from tensorflow.python.ops import nn_ops as nn_ops
from tensorflow.python.platform import googletest
from tensorflow.python.platform import test

@test_util.with_c_api
class IntegrationTest(test_util.TensorFlowTestCase):

  def setUp(self):
    """ Setup method """
    super(IntegrationTest, self).setUp()
    warnings.simplefilter('always')
    inp_dims = (100, 24, 24, 2)
    self._input = np.random.random_sample(inp_dims)
    self._original_graph = self.get_simple_graph_def()
    self._gpu_options = cpb2.GPUOptions(
        per_process_gpu_memory_fraction=0.50)
    self._config = cpb2.ConfigProto(gpu_options=self._gpu_options)
    self._reference = self.run_graph(self._original_graph, self._input)

  def get_simple_graph_def(self):
    """Create a simple graph and return its graph_def."""
    g = ops.Graph()
    with g.as_default():
      a = aops.placeholder(
          dtype=dtypes.float32, shape=(None, 24, 24, 2), name="input")
      e = cop.constant(
          [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
          name="weights",
          dtype=dtypes.float32)
      conv = nn.conv2d(
          input=a,
          filter=e,
          strides=[1, 2, 2, 1],
          padding="SAME",
          name="conv")
      b = cop.constant(
          [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtypes.float32)
      t = nn.bias_add(conv, b, name="biasAdd")
      relu = nn.relu(t, "relu")
      idty = aops.identity(relu, "ID")
      v = nn_ops.max_pool(
          idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
      aops.squeeze(v, name="output")
    return g.as_graph_def()

  def run_graph(self, gdef, dumm_inp):
    """Run given graphdef once."""
    ops.reset_default_graph()
    g = ops.Graph()
    with g.as_default():
      inp, out = importer.import_graph_def(
          graph_def=gdef, return_elements=["input", "output"])
      inp = inp.outputs[0]
      out = out.outputs[0]
    with self.test_session(
        graph=g, config=self._config, use_gpu=True,
        force_gpu=True) as sess:
      val = sess.run(out, {inp: dumm_inp})
    return val

  # Use real data that is representative of the inference dataset
  # for calibration. For this test script it is random data.
  def run_calibration(self, gdef, dumm_inp):
    """Run given calibration graph multiple times."""
    ops.reset_default_graph()
    g = ops.Graph()
    with g.as_default():
      inp, out = importer.import_graph_def(
          graph_def=gdef, return_elements=["input", "output"])
      inp = inp.outputs[0]
      out = out.outputs[0]
      # run over real calibration data here, we are mimicking a calibration set of
      # 30 different batches. Use as much calibration data as you want
    with self.test_session(
        graph=g, config=self._config, use_gpu=True,
        force_gpu=True) as sess:
      for _ in range(30):
        val = sess.run(out, {inp: dumm_inp})
    return val

  def get_trt_graph(self, mode):
    """  return trt converted graph """
    if mode == "FP32":
      return trt.create_inference_graph(
          input_graph_def=self._original_graph,
          outputs=["output"],
          max_batch_size=self._input.shape[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "FP32",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )
    elif mode == "FP16":
      return trt.create_inference_graph(
          input_graph_def=self._original_graph,
          outputs=["output"],
          max_batch_size=self._input.shape[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "FP16",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )
    elif mode == "INT8":
      return trt.create_inference_graph(
          input_graph_def=self._original_graph,
          outputs=["output"],
          max_batch_size=self._input.shape[0],
          max_workspace_size_bytes=1 << 25,
          precision_mode=
          "INT8",  # TRT Engine precision "FP32","FP16" or "INT8"
          minimum_segment_size=2  # minimum number of nodes in an engine
          )

    return None

  def testFP32(self):
    """ Test FP32 conversion. Results should be identical to native case """
    trt_graph = self.get_trt_graph("FP32")
    result = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(self._reference, result)
    result1 = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(result1, result)

  def testFP16(self):
    """ Test FP16 conversion. Results may be different from native case """
    trt_graph = self.get_trt_graph("FP16")
    result = self.run_graph(trt_graph, self._input)
    self.assertAllClose(self._reference, result,rtol=1.e-03)
    result1 = self.run_graph(trt_graph, self._input)
    self.assertAllEqual(result1, result)

  def testINT8(self):
    """ Test INT8 conversion. Results may be different from native case """
    calib_graph = self.get_trt_graph("INT8")
    result = self.run_calibration(calib_graph, self._input)
    self.assertAllEqual(self._reference, result)
    int8_graph = trt.calib_graph_to_infer_graph(calib_graph)
    result = self.run_graph(int8_graph, self._input)
    self.assertAllClose(self._reference, result,rtol=1.e-03)
    result1 = self.run_graph(int8_graph, self._input)
    self.assertAllEqual(result1, result)

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


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

List of built-in functions: https://docs.python.org/3/library/functions.html
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import inspect

import six

from tensorflow.python.autograph.utils import py_func
from tensorflow.python.autograph.utils import tensors
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_parsing_ops
from tensorflow.python.ops import gen_string_ops
from tensorflow.python.ops import list_ops
from tensorflow.python.ops import math_ops

UNSPECIFIED = object()

def overload_of(f):
  if f in SUPPORTED_BUILTINS:
    return BUILTIN_FUINCTIONS_MAP[f.__name__]
  return f

def _find_originating_frame(caller_fn_scope, innermost=True):
  """Locates the frame in which `caller_fn_scope` was defined."""
  ctx_frame = inspect.currentframe()
  result = None
  while ctx_frame is not None:
    # Note it should not be normally possible to get false positives this way
    # because the function scope object is not accessible to user code (barring
    # call stack introspection).
    if ctx_frame.f_locals.get(caller_fn_scope.name, None) is caller_fn_scope:
      result = ctx_frame
      if innermost:
        break
    ctx_frame = ctx_frame.f_back

  assert result is not None, (
      'the conversion process should ensure the caller_fn_scope is always'
      ' found somewhere on the call stack')

  return result

def eval_in_original_context(f, args, caller_fn_scope):
  """Executes the eval function in the context of a specified function."""
  # When control flow is rewritten using functions, eval should use the
  # variables found in the same block where it was called. That is equivalent
  # to the innermost function call.
  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=True)

  args = (
      args[0],
      ctx_frame.f_globals if len(args) < 2 else args[1],
      ctx_frame.f_locals if len(args) < 3 else args[2],
  )
  return f(*args)

def super_in_original_context(f, args, caller_fn_scope):
  """Executes the super function in the context of a specified function.

  See https://docs.python.org/3/library/functions.html#super for the exact
  details

  Args:
    f: Callable, typically the super builtin
    args: List[Any], the original call arguments
    caller_fn_scope: Optional[function_wrappers.FunctionScope], the function
      scope of the converted function in which this call was originally made

  Returns:
    The result of calling `f` as if it was called in the frame indicated by
      `caller_fn_scope`.
  """

  # Python 2 doesn't support implicit argument super variants.
  if six.PY2:
    return f(*args)

  # Only the no-arg call is desugared.
  if args:
    return f(*args)

  # Inner functions seem to include their closure in f_locals, so we need
  # to find the outermost frame.
  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=False)

  # When super(..) is called without arguments, it looks for __class__ cell
  # variable and the first argument passed in the enclosing function according
  # to the spec https://www.python.org/dev/peps/pep-3135/ .
  #
  # We couldn't verify if `inspect.currentframe().f_code.co_varnames[0]` is
  # guaranteed to be the first argument from an official doc or PEP, however,
  # it's fairly stable and well established:
  # - An unofficial community doc mentions it.
  #   https://python-reference.readthedocs.io/en/latest/docs/code/varnames.html
  # - CPython has tests checking that order, which was merged in 2008, and
  #   unchanged since then.
  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py2_test_grammar.py#L157
  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py3_test_grammar.py#L192
  #
  # Note: the name can be more reliably obtained by inspecting the calling
  # function's argspec.
  #
  # Even though methods can be declared using *args (def method(*args)),
  # that pattern is disallowed by super() -- it raises super() no arguments.
  # Method definitions using **kwargs are not allowed at all.
  # In other words, we can always assume that self is on the first positional
  # argument (for correct code).
  #
  # TODO(mdan): Consider additional checks in case the input code is incorrect.
  # For example, the error might be cryptic compared to what super() regularly
  # raises.

  type_arg = ctx_frame.f_locals['__class__']
  self_arg_name = ctx_frame.f_code.co_varnames[0]
  self_arg = ctx_frame.f_locals[self_arg_name]
  return f(type_arg, self_arg)

def abs_(x):
  if tensor_util.is_tensor(x):
    return _tf_abs(x)
  return _py_abs(x)

def _tf_abs(x):
  return math_ops.abs(x)

def _py_abs(x):
  return abs(x)

def float_(x=0):
  if tensor_util.is_tensor(x):
    return _tf_float(x)
  return _py_float(x)

def _tf_float(x):
  # TODO(mdan): We shouldn't assume float32.
  if x.dtype == dtypes.string:
    return gen_parsing_ops.string_to_number(x, out_type=dtypes.float32)
  return math_ops.cast(x, dtype=dtypes.float32)

def _py_float(x):
  return float(x)

def int_(x=0, base=UNSPECIFIED):
  if tensor_util.is_tensor(x):
    return _tf_int(x, base)
  return _py_int(x, base)

def _tf_int(x, base):
  if base not in (10, UNSPECIFIED):
    raise NotImplementedError('base {} not supported for int'.format(base))

  # TODO(mdan): We shouldn't assume int32.
  if x.dtype == dtypes.string:
    return gen_parsing_ops.string_to_number(x, out_type=dtypes.int32)
  return math_ops.cast(x, dtype=dtypes.int32)

def _py_int(x, base):
  if base is UNSPECIFIED:
    return int(x)
  return int(x, base)

def len_(s):
  if tensors.is_tensor_array(s):
    return _tf_tensor_array_len(s)
  elif tensors.is_tensor_list(s):
    return _tf_tensor_list_len(s)
  elif tensor_util.is_tensor(s):
    return _tf_tensor_len(s)
  return _py_len(s)

def _tf_tensor_array_len(s):
  return s.size()

def _tf_tensor_list_len(s):
  return list_ops.tensor_list_length(s)

def _tf_tensor_len(s):
  """Overload of len_ for Tensor arguments."""
  # Statically shaped tensors: length is known ahead of time.
  if s.shape.ndims and s.shape.dims[0].value is not None:
    return s.shape.dims[0].value

  # Static shape of unknown dimensions: use dynamic shape but statically
  # check that it's a scalar.
  shape = array_ops.shape(s)

  assert shape.shape, 'shape tensor of zero size? {}'.format(shape)

  if shape.shape[0] == 0:
    raise ValueError(
        'len requires a non-scalar tensor, got one of shape {}'.format(shape))

  if shape.shape.dims[0].value is not None:
    return array_ops.shape(s)[0]

  # Fully dynamic shape: use ops.
  rank = array_ops.rank(s)

  def raise_zero_rank_error():
    msg = gen_string_ops.string_join(
        ['len requires non-zero rank, got ',
         gen_string_ops.as_string(rank)])
    with ops.control_dependencies([control_flow_ops.Assert(False, [msg])]):
      return constant_op.constant(0, dtype=dtypes.int32)

  return control_flow_ops.cond(rank > 0, lambda: array_ops.shape(s)[0],
                               raise_zero_rank_error)

def _py_len(s):
  return len(s)

def print_(*objects, **kwargs):
  """Overload of the print builtin."""
  # Note: Python 2.6 doesn't support explicit keywords after starargs.
  unknown_kwargs = tuple(
      set(kwargs.keys()) - set(('sep', 'end', 'file', 'flush')))
  if unknown_kwargs:
    raise ValueError('invalid keyword arguments: {}'.format(unknown_kwargs))

  # TODO(mdan): Use next.flatten(objects) instead?
  if any(tensor_util.is_tensor(o) for o in objects):
    # TODO(mdan): use tf.print instead.
    return _tf_py_func_print(objects, kwargs)
  else:
    _py_print(*objects, **kwargs)

def _py_print(*objects, **kwargs):
  print(*objects, **kwargs)

def _tf_py_func_print(objects, kwargs):
  """Overload of print_ as a py_func implementation."""
  override_kwargs = {k: v for k, v in kwargs.items() if v is not UNSPECIFIED}
  if 'flush' not in override_kwargs:
    # Defaulting to flushing the console in graph mode, which helps reduce
    # garbled output in IPython.
    override_kwargs['flush'] = True

  def print_wrapper(*vals):
    vals = tuple(v.numpy() if tensor_util.is_tensor(v) else v for v in vals)
    if six.PY3:
      # TensorFlow doesn't seem to generate Unicode when passing strings to
      # py_func. This causes the print to add a "b'" wrapper to the output,
      # which is probably never what you want.
      vals = tuple(
          v.decode('utf-8') if isinstance(v, bytes) else v for v in vals)
    six.print_(*vals, **override_kwargs)

  return py_func.wrap_py_func(
      print_wrapper, None, objects, use_dummy_return=True)

def range_(start_or_stop, stop=UNSPECIFIED, step=UNSPECIFIED):
  if any(tensor_util.is_tensor(s) for s in (start_or_stop, stop, step)):
    return _tf_range(start_or_stop, stop, step)
  return _py_range(start_or_stop, stop, step)

def _tf_range(start_or_stop, stop, step):
  """Overload of range_ that generates a TF range tensor."""
  # Note: for static inputs (e.g. constants), tf.range errors out at graph
  # construction time, instead of returning an empty tensor. Preventing the
  # graph construction error aligns the semantics with Python.

  # TODO(mdan): We should optimize this when a full tensor is not required.
  if step is not UNSPECIFIED:
    # TODO(mdan): Add argument coercion similar to other cases.
    return math_ops.range(start_or_stop, stop, step)
  if stop is not UNSPECIFIED:
    stop = math_ops.maximum(start_or_stop, stop)
    return math_ops.range(start_or_stop, stop)
  start_or_stop = math_ops.maximum(start_or_stop, 0)
  return math_ops.range(start_or_stop)

def _py_range(start_or_stop, stop, step):
  if step is not UNSPECIFIED:
    return range(start_or_stop, stop, step)
  if stop is not UNSPECIFIED:
    return range(start_or_stop, stop)
  return range(start_or_stop)

def enumerate_(s, start=0):
  if isinstance(s, dataset_ops.DatasetV2):
    return _tf_dataset_enumerate(s, start)
  return _py_enumerate(s, start)

def _tf_dataset_enumerate(s, start=0):
  return s.enumerate(start)

def _py_enumerate(s, start=0):
  return enumerate(s, start)

def zip_(*s):
  if all(isinstance(x, dataset_ops.DatasetV2) for x in s):
    return _tf_dataset_zip(*s)
  return _py_zip(*s)

def _tf_dataset_zip(*s):
  tup = ()
  for x in s:
    tup += (x,)
  return dataset_ops.DatasetV2.zip(tup)

def _py_zip(*s):
  return zip(*s)

SUPPORTED_BUILTINS = (abs, float, int, len, print, range, enumerate, zip)

if six.PY2:
  SUPPORTED_BUILTINS += (xrange,)

BUILTIN_FUINCTIONS_MAP = {
    'abs': abs_,
    'float': float_,
    'int': int_,
    'len': len_,
    'print': print_,
    'range': range_,
    # TODO(mdan): This might make more sense as tf.data.range.
    'xrange': range_,
    'enumerate': enumerate_,
    'zip': zip_,
}
<commit_msg>Fix argument naming and tuple on zip<commit_after># Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Operators corresponding to Python builtin functions.

List of built-in functions: https://docs.python.org/3/library/functions.html
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import inspect

import six

from tensorflow.python.autograph.utils import py_func
from tensorflow.python.autograph.utils import tensors
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_parsing_ops
from tensorflow.python.ops import gen_string_ops
from tensorflow.python.ops import list_ops
from tensorflow.python.ops import math_ops

UNSPECIFIED = object()

def overload_of(f):
  if f in SUPPORTED_BUILTINS:
    return BUILTIN_FUINCTIONS_MAP[f.__name__]
  return f

def _find_originating_frame(caller_fn_scope, innermost=True):
  """Locates the frame in which `caller_fn_scope` was defined."""
  ctx_frame = inspect.currentframe()
  result = None
  while ctx_frame is not None:
    # Note it should not be normally possible to get false positives this way
    # because the function scope object is not accessible to user code (barring
    # call stack introspection).
    if ctx_frame.f_locals.get(caller_fn_scope.name, None) is caller_fn_scope:
      result = ctx_frame
      if innermost:
        break
    ctx_frame = ctx_frame.f_back

  assert result is not None, (
      'the conversion process should ensure the caller_fn_scope is always'
      ' found somewhere on the call stack')

  return result

def eval_in_original_context(f, args, caller_fn_scope):
  """Executes the eval function in the context of a specified function."""
  # When control flow is rewritten using functions, eval should use the
  # variables found in the same block where it was called. That is equivalent
  # to the innermost function call.
  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=True)

  args = (
      args[0],
      ctx_frame.f_globals if len(args) < 2 else args[1],
      ctx_frame.f_locals if len(args) < 3 else args[2],
  )
  return f(*args)

def super_in_original_context(f, args, caller_fn_scope):
  """Executes the super function in the context of a specified function.

  See https://docs.python.org/3/library/functions.html#super for the exact
  details

  Args:
    f: Callable, typically the super builtin
    args: List[Any], the original call arguments
    caller_fn_scope: Optional[function_wrappers.FunctionScope], the function
      scope of the converted function in which this call was originally made

  Returns:
    The result of calling `f` as if it was called in the frame indicated by
      `caller_fn_scope`.
  """

  # Python 2 doesn't support implicit argument super variants.
  if six.PY2:
    return f(*args)

  # Only the no-arg call is desugared.
  if args:
    return f(*args)

  # Inner functions seem to include their closure in f_locals, so we need
  # to find the outermost frame.
  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=False)

  # When super(..) is called without arguments, it looks for __class__ cell
  # variable and the first argument passed in the enclosing function according
  # to the spec https://www.python.org/dev/peps/pep-3135/ .
  #
  # We couldn't verify if `inspect.currentframe().f_code.co_varnames[0]` is
  # guaranteed to be the first argument from an official doc or PEP, however,
  # it's fairly stable and well established:
  # - An unofficial community doc mentions it.
  #   https://python-reference.readthedocs.io/en/latest/docs/code/varnames.html
  # - CPython has tests checking that order, which was merged in 2008, and
  #   unchanged since then.
  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py2_test_grammar.py#L157
  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py3_test_grammar.py#L192
  #
  # Note: the name can be more reliably obtained by inspecting the calling
  # function's argspec.
  #
  # Even though methods can be declared using *args (def method(*args)),
  # that pattern is disallowed by super() -- it raises super() no arguments.
  # Method definitions using **kwargs are not allowed at all.
  # In other words, we can always assume that self is on the first positional
  # argument (for correct code).
  #
  # TODO(mdan): Consider additional checks in case the input code is incorrect.
  # For example, the error might be cryptic compared to what super() regularly
  # raises.

  type_arg = ctx_frame.f_locals['__class__']
  self_arg_name = ctx_frame.f_code.co_varnames[0]
  self_arg = ctx_frame.f_locals[self_arg_name]
  return f(type_arg, self_arg)

def abs_(x):
  if tensor_util.is_tensor(x):
    return _tf_abs(x)
  return _py_abs(x)

def _tf_abs(x):
  return math_ops.abs(x)

def _py_abs(x):
  return abs(x)

def float_(x=0):
  if tensor_util.is_tensor(x):
    return _tf_float(x)
  return _py_float(x)

def _tf_float(x):
  # TODO(mdan): We shouldn't assume float32.
  if x.dtype == dtypes.string:
    return gen_parsing_ops.string_to_number(x, out_type=dtypes.float32)
  return math_ops.cast(x, dtype=dtypes.float32)

def _py_float(x):
  return float(x)

def int_(x=0, base=UNSPECIFIED):
  if tensor_util.is_tensor(x):
    return _tf_int(x, base)
  return _py_int(x, base)

def _tf_int(x, base):
  if base not in (10, UNSPECIFIED):
    raise NotImplementedError('base {} not supported for int'.format(base))

  # TODO(mdan): We shouldn't assume int32.
  if x.dtype == dtypes.string:
    return gen_parsing_ops.string_to_number(x, out_type=dtypes.int32)
  return math_ops.cast(x, dtype=dtypes.int32)

def _py_int(x, base):
  if base is UNSPECIFIED:
    return int(x)
  return int(x, base)

def len_(s):
  if tensors.is_tensor_array(s):
    return _tf_tensor_array_len(s)
  elif tensors.is_tensor_list(s):
    return _tf_tensor_list_len(s)
  elif tensor_util.is_tensor(s):
    return _tf_tensor_len(s)
  return _py_len(s)

def _tf_tensor_array_len(s):
  return s.size()

def _tf_tensor_list_len(s):
  return list_ops.tensor_list_length(s)

def _tf_tensor_len(s):
  """Overload of len_ for Tensor arguments."""
  # Statically shaped tensors: length is known ahead of time.
  if s.shape.ndims and s.shape.dims[0].value is not None:
    return s.shape.dims[0].value

  # Static shape of unknown dimensions: use dynamic shape but statically
  # check that it's a scalar.
  shape = array_ops.shape(s)

  assert shape.shape, 'shape tensor of zero size? {}'.format(shape)

  if shape.shape[0] == 0:
    raise ValueError(
        'len requires a non-scalar tensor, got one of shape {}'.format(shape))

  if shape.shape.dims[0].value is not None:
    return array_ops.shape(s)[0]

  # Fully dynamic shape: use ops.
  rank = array_ops.rank(s)

  def raise_zero_rank_error():
    msg = gen_string_ops.string_join(
        ['len requires non-zero rank, got ',
         gen_string_ops.as_string(rank)])
    with ops.control_dependencies([control_flow_ops.Assert(False, [msg])]):
      return constant_op.constant(0, dtype=dtypes.int32)

  return control_flow_ops.cond(rank > 0, lambda: array_ops.shape(s)[0],
                               raise_zero_rank_error)

def _py_len(s):
  return len(s)

def print_(*objects, **kwargs):
  """Overload of the print builtin."""
  # Note: Python 2.6 doesn't support explicit keywords after starargs.
  unknown_kwargs = tuple(
      set(kwargs.keys()) - set(('sep', 'end', 'file', 'flush')))
  if unknown_kwargs:
    raise ValueError('invalid keyword arguments: {}'.format(unknown_kwargs))

  # TODO(mdan): Use next.flatten(objects) instead?
  if any(tensor_util.is_tensor(o) for o in objects):
    # TODO(mdan): use tf.print instead.
    return _tf_py_func_print(objects, kwargs)
  else:
    _py_print(*objects, **kwargs)

def _py_print(*objects, **kwargs):
  print(*objects, **kwargs)

def _tf_py_func_print(objects, kwargs):
  """Overload of print_ as a py_func implementation."""
  override_kwargs = {k: v for k, v in kwargs.items() if v is not UNSPECIFIED}
  if 'flush' not in override_kwargs:
    # Defaulting to flushing the console in graph mode, which helps reduce
    # garbled output in IPython.
    override_kwargs['flush'] = True

  def print_wrapper(*vals):
    vals = tuple(v.numpy() if tensor_util.is_tensor(v) else v for v in vals)
    if six.PY3:
      # TensorFlow doesn't seem to generate Unicode when passing strings to
      # py_func. This causes the print to add a "b'" wrapper to the output,
      # which is probably never what you want.
      vals = tuple(
          v.decode('utf-8') if isinstance(v, bytes) else v for v in vals)
    six.print_(*vals, **override_kwargs)

  return py_func.wrap_py_func(
      print_wrapper, None, objects, use_dummy_return=True)

def range_(start_or_stop, stop=UNSPECIFIED, step=UNSPECIFIED):
  if any(tensor_util.is_tensor(s) for s in (start_or_stop, stop, step)):
    return _tf_range(start_or_stop, stop, step)
  return _py_range(start_or_stop, stop, step)

def _tf_range(start_or_stop, stop, step):
  """Overload of range_ that generates a TF range tensor."""
  # Note: for static inputs (e.g. constants), tf.range errors out at graph
  # construction time, instead of returning an empty tensor. Preventing the
  # graph construction error aligns the semantics with Python.

  # TODO(mdan): We should optimize this when a full tensor is not required.
  if step is not UNSPECIFIED:
    # TODO(mdan): Add argument coercion similar to other cases.
    return math_ops.range(start_or_stop, stop, step)
  if stop is not UNSPECIFIED:
    stop = math_ops.maximum(start_or_stop, stop)
    return math_ops.range(start_or_stop, stop)
  start_or_stop = math_ops.maximum(start_or_stop, 0)
  return math_ops.range(start_or_stop)

def _py_range(start_or_stop, stop, step):
  if step is not UNSPECIFIED:
    return range(start_or_stop, stop, step)
  if stop is not UNSPECIFIED:
    return range(start_or_stop, stop)
  return range(start_or_stop)

def enumerate_(s, start=0):
  if isinstance(s, dataset_ops.DatasetV2):
    return _tf_dataset_enumerate(s, start)
  return _py_enumerate(s, start)

def _tf_dataset_enumerate(s, start=0):
  return s.enumerate(start)

def _py_enumerate(s, start=0):
  return enumerate(s, start)

def zip_(*iterables):
  if all(isinstance(x, dataset_ops.DatasetV2) for x in iterables):
    return _tf_dataset_zip(*iterables)
  return _py_zip(*iterables)

def _tf_dataset_zip(*iterables):
  return dataset_ops.DatasetV2.zip(tuple(iterables))

def _py_zip(*iterables):
  return zip(*iterables)

SUPPORTED_BUILTINS = (abs, float, int, len, print, range, enumerate, zip)

if six.PY2:
  SUPPORTED_BUILTINS += (xrange,)

BUILTIN_FUINCTIONS_MAP = {
    'abs': abs_,
    'float': float_,
    'int': int_,
    'len': len_,
    'print': print_,
    'range': range_,
    # TODO(mdan): This might make more sense as tf.data.range.
    'xrange': range_,
    'enumerate': enumerate_,
    'zip': zip_,
}
<|endoftext|>


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

"""Templates to generate dataset builders.

Depending on the given format, it will generate code to add format-specific
dataset builders to tfds.
"""
import textwrap

from tensorflow_datasets.scripts.cli import cli_utils

# Names of dataset builders.
STANDARD = 'standard'
CONLL = 'conll'
CONLLU = 'conllu'

def create_builder_template(info: cli_utils.DatasetInfo):
  """Calls the required (possibly format-specific) dataset builder template.

  Args:
    info: `cli_utils.DatasetInfo` object containing all information about the
      dataset necessary to generate a template for the new dataset.

  Returns:
    A string containing the builder-specific template for the new dataset, to be
    filled by the user.
  """
  if info.data_format == STANDARD:
    return _standard_template(info)
  elif info.data_format == CONLL:
    return _conll_template(info)
  elif info.data_format == CONLLU:
    return _conllu_template(info)
  else:
    raise ValueError(
        f'Required format {info.data_format} isn\'t associated with '
        'a format-specific builder in TFDS.')

def _standard_template(info: cli_utils.DatasetInfo) -> str:
  """Returns a template for a `tfds.core.GeneratorBasedBuilder`."""
  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.core.GeneratorBasedBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the tfds.core.DatasetInfo object
          return tfds.core.DatasetInfo(
              builder=self,
              description=_DESCRIPTION,
              features=tfds.features.FeaturesDict({{
                  # These are the features of your dataset like images, labels ...
                  'image': tfds.features.Image(shape=(None, None, 3)),
                  'label': tfds.features.ClassLabel(names=['no', 'yes']),
              }}),
              # If there's a common (input, target) tuple from the
              # features, specify them here. They'll be used if
              # `as_supervised=True` in `builder.as_dataset`.
              supervised_keys=('image', 'label'),  # Set to `None` to disable
              homepage='https://dataset-homepage/',
              citation=_CITATION,
          )

        def _split_generators(self, dl_manager: tfds.download.DownloadManager):
          """Returns SplitGenerators."""
          # {info.todo}: Downloads the data and defines the splits
          path = dl_manager.download_and_extract('https://todo-data-url')

          # {info.todo}: Returns the Dict[split names, Iterator[Key, Example]]
          return {{
              'train': self._generate_examples(path / 'train_imgs'),
          }}

        def _generate_examples(self, path):
          """Yields examples."""
          # {info.todo}: Yields (key, example) tuples from the dataset
          for f in path.glob('*.jpeg'):
            yield 'key', {{
                'image': f,
                'label': 'yes',
            }}
      ''')
  return content

def _conll_template(info: cli_utils.DatasetInfo) -> str:
  """A template for ConllDatasetBuilder."""

  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      from tensorflow_datasets.core.dataset_builders.conll import conll_dataset_builder_utils as conll_lib
      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.dataset_builders.ConllDatasetBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}
        # {info.todo}: Add details about the dataset's features.
        # conll_lib contains a set of ready-to-use features.
        BUILDER_CONFIGS = [conll_lib.CONLL_2002_CONFIG]

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the dataset infos.
          return self.create_dataset_info(
            description=_DESCRIPTION,
            homepage="",
            citation=_CITATION,
          )

        def _split_generators(self, dl_manager: tfds.download.DownloadManager):
          """Returns SplitGenerators."""
          # {info.todo}: Downloads the data and defines the splits
          path = dl_manager.download_and_extract('https://todo-data-url')

          return {{
            'train': self._generate_examples(path / 'train.txt',)
          }}

        # {info.todo}: If you need a customized _generate_examples function,
        # comment out the following.
        # def _generate_examples(
        #     self,
        #     path: Union[epath.PathLike, List[epath.PathLike]],
        #   ):
        #   """Yields (key, example) examples."""
        #   pass
      ''')
  return content

def _conllu_template(info: cli_utils.DatasetInfo) -> str:
  """A template for ConllUDatasetBuilder."""

  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      from tensorflow_datasets.core.dataset_builders.conll import conllu_dataset_builder_utils as conllu_lib
      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.dataset_builders.ConllUDatasetBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}
        # {info.todo}: Add details about the dataset's config to use to
        # BUILDER_CONFIGS. conllu_lib contains a set of ready-to-use features.
        BUILDER_CONFIGS = [
            conllu_lib.get_universal_morphology_config(
                language='',
                features=conllu_lib.UNIVERSAL_DEPENDENCIES_FEATURES,
                name='')
        ]

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the dataset infos.
          return self.create_dataset_info(
            description=_DESCRIPTION,
            homepage='',
            citation=_CITATION,
          )

        # {info.todo}: Specify the process_example_fn to be used in the
        # `_generate_examples` method to process examples.
        # The default `process_example_fn` processes a conllu-annotated example
        # using the features specified in BUILDER_CONFIGS.
        # `conllu_dataset_builder` already provides a number of ready-to-use
        # process functions.
        return {{
            'train':
                self._generate_examples(
                    filepaths=path / 'train.txt',
                    # Remove if you want to use the default `process_example_fn`.
                    # process_example_fn=
                )
        }}

        # {info.todo}: If you need a customized _generate_examples function,
        # comment out the following, otherwise remove it.
        # def _generate_examples(
        #     self,
        #     path: Union[epath.PathLike, List[epath.PathLike]],
        #   ):
        #   """Yields (key, example) examples."""
        #   pass
      ''')
  return content
<commit_msg>Fix bug in CoNLL-U CLI<commit_after># coding=utf-8
# Copyright 2022 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Templates to generate dataset builders.

Depending on the given format, it will generate code to add format-specific
dataset builders to tfds.
"""
import textwrap

from tensorflow_datasets.scripts.cli import cli_utils

# Names of dataset builders.
STANDARD = 'standard'
CONLL = 'conll'
CONLLU = 'conllu'

def create_builder_template(info: cli_utils.DatasetInfo):
  """Calls the required (possibly format-specific) dataset builder template.

  Args:
    info: `cli_utils.DatasetInfo` object containing all information about the
      dataset necessary to generate a template for the new dataset.

  Returns:
    A string containing the builder-specific template for the new dataset, to be
    filled by the user.
  """
  if info.data_format == STANDARD:
    return _standard_template(info)
  elif info.data_format == CONLL:
    return _conll_template(info)
  elif info.data_format == CONLLU:
    return _conllu_template(info)
  else:
    raise ValueError(
        f'Required format {info.data_format} isn\'t associated with '
        'a format-specific builder in TFDS.')

def _standard_template(info: cli_utils.DatasetInfo) -> str:
  """Returns a template for a `tfds.core.GeneratorBasedBuilder`."""
  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.core.GeneratorBasedBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the tfds.core.DatasetInfo object
          return tfds.core.DatasetInfo(
              builder=self,
              description=_DESCRIPTION,
              features=tfds.features.FeaturesDict({{
                  # These are the features of your dataset like images, labels ...
                  'image': tfds.features.Image(shape=(None, None, 3)),
                  'label': tfds.features.ClassLabel(names=['no', 'yes']),
              }}),
              # If there's a common (input, target) tuple from the
              # features, specify them here. They'll be used if
              # `as_supervised=True` in `builder.as_dataset`.
              supervised_keys=('image', 'label'),  # Set to `None` to disable
              homepage='https://dataset-homepage/',
              citation=_CITATION,
          )

        def _split_generators(self, dl_manager: tfds.download.DownloadManager):
          """Returns SplitGenerators."""
          # {info.todo}: Downloads the data and defines the splits
          path = dl_manager.download_and_extract('https://todo-data-url')

          # {info.todo}: Returns the Dict[split names, Iterator[Key, Example]]
          return {{
              'train': self._generate_examples(path / 'train_imgs'),
          }}

        def _generate_examples(self, path):
          """Yields examples."""
          # {info.todo}: Yields (key, example) tuples from the dataset
          for f in path.glob('*.jpeg'):
            yield 'key', {{
                'image': f,
                'label': 'yes',
            }}
      ''')
  return content

def _conll_template(info: cli_utils.DatasetInfo) -> str:
  """A template for ConllDatasetBuilder."""

  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      from tensorflow_datasets.core.dataset_builders.conll import conll_dataset_builder_utils as conll_lib
      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.dataset_builders.ConllDatasetBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}
        # {info.todo}: Add details about the dataset's features.
        # conll_lib contains a set of ready-to-use features.
        BUILDER_CONFIGS = [conll_lib.CONLL_2002_CONFIG]

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the dataset infos.
          return self.create_dataset_info(
            description=_DESCRIPTION,
            homepage="",
            citation=_CITATION,
          )

        def _split_generators(self, dl_manager: tfds.download.DownloadManager):
          """Returns SplitGenerators."""
          # {info.todo}: Downloads the data and defines the splits
          path = dl_manager.download_and_extract('https://todo-data-url')

          return {{
            'train': self._generate_examples(path / 'train.txt',)
          }}

        # {info.todo}: If you need a customized _generate_examples function,
        # comment out the following.
        # def _generate_examples(
        #     self,
        #     path: Union[epath.PathLike, List[epath.PathLike]],
        #   ):
        #   """Yields (key, example) examples."""
        #   pass
      ''')
  return content

def _conllu_template(info: cli_utils.DatasetInfo) -> str:
  """A template for ConllUDatasetBuilder."""

  content = textwrap.dedent(f'''\
      """{info.name} dataset."""

      from tensorflow_datasets.core.dataset_builders.conll import conllu_dataset_builder_utils as conllu_lib
      import {info.tfds_api} as tfds

      # {info.todo}: Markdown description  that will appear on the catalog page.
      _DESCRIPTION = """
      Description is **formatted** as markdown.

      It should also contain any processing which has been applied (if any),
      (e.g. corrupted example skipped, images cropped,...):
      """

      # {info.todo}: BibTeX citation
      _CITATION = """
      """

      class {info.cls_name}(tfds.dataset_builders.ConllUDatasetBuilder):
        """DatasetBuilder for {info.name} dataset."""

        VERSION = tfds.core.Version('1.0.0')
        RELEASE_NOTES = {{
            '1.0.0': 'Initial release.',
        }}
        # {info.todo}: Add details about the dataset's config to use to
        # BUILDER_CONFIGS. conllu_lib contains a set of ready-to-use features.
        BUILDER_CONFIGS = [
            conllu_lib.get_universal_morphology_config(
                language='',
                features=conllu_lib.UNIVERSAL_DEPENDENCIES_FEATURES,
                name='')
        ]

        def _info(self) -> tfds.core.DatasetInfo:
          """Returns the dataset metadata."""
          # {info.todo}: Specifies the dataset infos.
          return self.create_dataset_info(
              description=_DESCRIPTION,
              homepage='',
              citation=_CITATION,
          )

        def _split_generators(self, dl_manager: tfds.download.DownloadManager):
          """Returns SplitGenerators."""
          # {info.todo}: Downloads the data and defines the splits
          path = dl_manager.download_and_extract('https://todo-data-url')

          # {info.todo}: Specify the process_example_fn to be used in the
          # `_generate_examples` method to process examples.
          # The default `process_example_fn` processes a conllu-annotated
          # example using the features specified in BUILDER_CONFIGS.
          # `conllu_dataset_builder` already provides a number of ready-to-use
          # process functions.
          return {{
              'train':
                  self._generate_examples(
                      filepaths=path / 'train.txt',
                      # Remove if you want to use the default `process_example_fn`.
                      # process_example_fn=
                  )
          }}

        # {info.todo}: If you need a customized _generate_examples function,
        # comment out the following, otherwise remove it.
        # def _generate_examples(
        #     self,
        #     path: Union[epath.PathLike, List[epath.PathLike]],
        #   ):
        #   """Yields (key, example) examples."""
        #   pass
      ''')
  return content
<|endoftext|>