<commit_before>import ctypes
import pkg_resources
import os
import sys
from typing import List, Optional

from setuptools import setup

VERSION = '10.4.0'

# List of packages supported by this version of CuPy.
PACKAGES = [
    'cupy-cuda102',
    'cupy-cuda110',
    'cupy-cuda111',
    'cupy-cuda112',
    'cupy-cuda113',
    'cupy-cuda114',
    'cupy-cuda115',
    'cupy-cuda116',
    'cupy-cuda117',
    'cupy-rocm-4-0',
    'cupy-rocm-4-2',
    'cupy-rocm-4-3',
    'cupy-rocm-5-0',
]

# List of packages NOT supported by this version of CuPy.
PACKAGES_OUTDATED = [
    'cupy-cuda80',
    'cupy-cuda90',
    'cupy-cuda91',
    'cupy-cuda92',
    'cupy-cuda100',
    'cupy-cuda101',
]

# List of sdist packages.
PACKAGES_SDIST = [
    'cupy',
]

class AutoDetectionFailed(Exception):
    pass

def _log(msg: str) -> None:
    sys.stdout.write(f'[cupy-wheel] {msg}\n')
    sys.stdout.flush()

def _get_version_from_library(
        libnames: List[str], funcname: str) -> Optional[int]:
    """Returns the library version from list of candidate libraries."""

    for libname in libnames:
        try:
            _log(f'Looking for library: {libname}')
            runtime_so = ctypes.CDLL(libname)
            break
        except Exception as e:
            _log(f'Failed to open {libname}: {e}')
    else:
        _log('No more candidate library to find')
        return None

    func = getattr(runtime_so, funcname, None)
    if func is None:
        raise AutoDetectionFailed(
            f'{libname}: {func} could not be found')
    func.restype = ctypes.c_int  # nvrtcResult
    func.argtypes = [
        ctypes.POINTER(ctypes.c_int),
        ctypes.POINTER(ctypes.c_int),
    ]
    major = ctypes.c_int()
    minor = ctypes.c_int()
    retval = func(major, minor)
    if retval != 0:  # NVRTC_SUCCESS
        raise AutoDetectionFailed(
            f'{libname}: {func} returned error: {retval}')
    version = major.value * 1000 + minor.value * 10
    _log(f'Detected version: {version}')
    return version

def _setup_win32_dll_directory() -> None:
    if not hasattr(os, 'add_dll_directory'):
        # Python 3.7 or earlier.
        return
    cuda_path = os.environ.get('CUDA_PATH', None)
    if cuda_path is None:
        _log('CUDA_PATH is not set.'
             'cupy-wheel may not be able to discover NVRTC to probe version')
        return
    os.add_dll_directory(os.path.join(cuda_path, 'bin'))

def _get_cuda_version() -> Optional[int]:
    """Returns the detected CUDA version or None."""

    if sys.platform == 'linux':
        libnames = [
            'libnvrtc.so.11.2',
            'libnvrtc.so.11.1',
            'libnvrtc.so.11.0',
            'libnvrtc.so.10.2',
        ]
    elif sys.platform == 'win32':
        libnames = [
            'nvrtc64_112_0.dll',
            'nvrtc64_111_0.dll',
            'nvrtc64_110_0.dll',
            'nvrtc64_102_0.dll',
        ]
        _setup_win32_dll_directory()
    else:
        _log(f'CUDA detection unsupported on platform: {sys.platform}')
        return None
    _log(f'Trying to detect CUDA version from libraries: {libnames}')
    version = _get_version_from_library(libnames, 'nvrtcVersion')
    return version

def _get_rocm_version() -> Optional[int]:
    """Returns the detected ROCm version or None."""
    if sys.platform == 'linux':
        libnames = ['libamdhip64.so']
    else:
        _log(f'ROCm detection unsupported on platform: {sys.platform}')
        return None
    version = _get_version_from_library(libnames, 'hiprtcVersion')
    return version

def _find_installed_packages() -> List[str]:
    """Returns the list of CuPy packages installed in the environment."""

    found = []
    for pkg in (PACKAGES + PACKAGES_OUTDATED + PACKAGES_SDIST):
        try:
            pkg_resources.get_distribution(pkg)
            found.append(pkg)
        except pkg_resources.DistributionNotFound:
            pass
    return found

def _cuda_version_to_package(ver: int) -> str:
    if ver < 10020:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too old.')
    elif ver < 11000:
        # CUDA 10.2
        suffix = '102'
    elif ver < 11010:
        # CUDA 11.0
        suffix = '110'
    elif ver < 11020:
        # CUDA 11.1
        suffix = '111'
    elif ver < 11030:
        # CUDA 11.2
        suffix = '112'
    elif ver < 11040:
        # CUDA 11.3
        suffix = '113'
    elif ver < 11050:
        # CUDA 11.4
        suffix = '114'
    elif ver < 11060:
        # CUDA 11.5
        suffix = '115'
    elif ver < 11070:
        # CUDA 11.6
        suffix = '116'
    elif ver < 11080:
        # CUDA 11.7
        suffix = '117'
    else:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too new.')
    return f'cupy-cuda{suffix}'

def _rocm_version_to_package(ver: int) -> str:
    if 400 <= ver < 410:
        # ROCm 4.0
        suffix = '4-0'
    elif 420 <= ver < 430:
        # ROCm 4.2
        suffix = '4-2'
    elif 4_03_00000 <= ver < 4_04_00000:
        # ROCm 4.3
        suffix = '4-3'
    elif 5_00_00000 <= ver < 5_01_00000:
        # ROCm 5.0
        suffix = '5-0'
    else:
        raise AutoDetectionFailed(
            f'Your ROCm version ({ver}) is unsupported.')
    return f'cupy-rocm-{suffix}'

def infer_best_package() -> str:
    """Returns the appropriate CuPy wheel package name for the environment."""

    # Find the existing CuPy wheel installation.
    installed = _find_installed_packages()
    if 1 < len(installed):
        raise AutoDetectionFailed(
            'You have multiple CuPy packages installed: \n'
            f'  {installed}\n'
            'Please uninstall all of them first, then try reinstalling.')

    elif 1 == len(installed):
        if installed[0] in PACKAGES_SDIST:
            raise AutoDetectionFailed(
                'You already have CuPy installed via source'
                ' (pip install cupy).')
        if installed[0] in PACKAGES_OUTDATED:
            raise AutoDetectionFailed(
                f'You have CuPy package "{installed[0]}" installed, but the'
                f' package is not available for version {VERSION}.')
        return installed[0]

    # Try CUDA.
    version = _get_cuda_version()
    if version is not None:
        return _cuda_version_to_package(version)

    # Try ROCm.
    version = _get_rocm_version()
    if version is not None:
        return _rocm_version_to_package(version)

    raise AutoDetectionFailed(
        'Unable to detect NVIDIA CUDA or AMD ROCm installation.')

#
# Entrypoint
#

def main() -> None:
    if os.environ.get('CUPY_UNIVERSAL_PKG_BUILD', None) is None:
        package = infer_best_package()
        requires = f'{package}=={VERSION}'
        _log(f'Installing package: {requires}')
        install_requires = [requires]
    else:
        _log('Building cupy-wheel package for release.')
        install_requires = []

    setup(
        name='cupy-wheel',
        version=VERSION,
        install_requires=install_requires,
    )

if __name__ == '__main__':
    main()
<commit_msg>type ignore as this only exist on Windows<commit_after>import ctypes
import pkg_resources
import os
import sys
from typing import List, Optional

from setuptools import setup

VERSION = '10.4.0'

# List of packages supported by this version of CuPy.
PACKAGES = [
    'cupy-cuda102',
    'cupy-cuda110',
    'cupy-cuda111',
    'cupy-cuda112',
    'cupy-cuda113',
    'cupy-cuda114',
    'cupy-cuda115',
    'cupy-cuda116',
    'cupy-cuda117',
    'cupy-rocm-4-0',
    'cupy-rocm-4-2',
    'cupy-rocm-4-3',
    'cupy-rocm-5-0',
]

# List of packages NOT supported by this version of CuPy.
PACKAGES_OUTDATED = [
    'cupy-cuda80',
    'cupy-cuda90',
    'cupy-cuda91',
    'cupy-cuda92',
    'cupy-cuda100',
    'cupy-cuda101',
]

# List of sdist packages.
PACKAGES_SDIST = [
    'cupy',
]

class AutoDetectionFailed(Exception):
    pass

def _log(msg: str) -> None:
    sys.stdout.write(f'[cupy-wheel] {msg}\n')
    sys.stdout.flush()

def _get_version_from_library(
        libnames: List[str], funcname: str) -> Optional[int]:
    """Returns the library version from list of candidate libraries."""

    for libname in libnames:
        try:
            _log(f'Looking for library: {libname}')
            runtime_so = ctypes.CDLL(libname)
            break
        except Exception as e:
            _log(f'Failed to open {libname}: {e}')
    else:
        _log('No more candidate library to find')
        return None

    func = getattr(runtime_so, funcname, None)
    if func is None:
        raise AutoDetectionFailed(
            f'{libname}: {func} could not be found')
    func.restype = ctypes.c_int  # nvrtcResult
    func.argtypes = [
        ctypes.POINTER(ctypes.c_int),
        ctypes.POINTER(ctypes.c_int),
    ]
    major = ctypes.c_int()
    minor = ctypes.c_int()
    retval = func(major, minor)
    if retval != 0:  # NVRTC_SUCCESS
        raise AutoDetectionFailed(
            f'{libname}: {func} returned error: {retval}')
    version = major.value * 1000 + minor.value * 10
    _log(f'Detected version: {version}')
    return version

def _setup_win32_dll_directory() -> None:
    if not hasattr(os, 'add_dll_directory'):
        # Python 3.7 or earlier.
        return
    cuda_path = os.environ.get('CUDA_PATH', None)
    if cuda_path is None:
        _log('CUDA_PATH is not set.'
             'cupy-wheel may not be able to discover NVRTC to probe version')
        return
    os.add_dll_directory(os.path.join(cuda_path, 'bin'))  # type: ignore[attr-defined] # NOQA

def _get_cuda_version() -> Optional[int]:
    """Returns the detected CUDA version or None."""

    if sys.platform == 'linux':
        libnames = [
            'libnvrtc.so.11.2',
            'libnvrtc.so.11.1',
            'libnvrtc.so.11.0',
            'libnvrtc.so.10.2',
        ]
    elif sys.platform == 'win32':
        libnames = [
            'nvrtc64_112_0.dll',
            'nvrtc64_111_0.dll',
            'nvrtc64_110_0.dll',
            'nvrtc64_102_0.dll',
        ]
        _setup_win32_dll_directory()
    else:
        _log(f'CUDA detection unsupported on platform: {sys.platform}')
        return None
    _log(f'Trying to detect CUDA version from libraries: {libnames}')
    version = _get_version_from_library(libnames, 'nvrtcVersion')
    return version

def _get_rocm_version() -> Optional[int]:
    """Returns the detected ROCm version or None."""
    if sys.platform == 'linux':
        libnames = ['libamdhip64.so']
    else:
        _log(f'ROCm detection unsupported on platform: {sys.platform}')
        return None
    version = _get_version_from_library(libnames, 'hiprtcVersion')
    return version

def _find_installed_packages() -> List[str]:
    """Returns the list of CuPy packages installed in the environment."""

    found = []
    for pkg in (PACKAGES + PACKAGES_OUTDATED + PACKAGES_SDIST):
        try:
            pkg_resources.get_distribution(pkg)
            found.append(pkg)
        except pkg_resources.DistributionNotFound:
            pass
    return found

def _cuda_version_to_package(ver: int) -> str:
    if ver < 10020:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too old.')
    elif ver < 11000:
        # CUDA 10.2
        suffix = '102'
    elif ver < 11010:
        # CUDA 11.0
        suffix = '110'
    elif ver < 11020:
        # CUDA 11.1
        suffix = '111'
    elif ver < 11030:
        # CUDA 11.2
        suffix = '112'
    elif ver < 11040:
        # CUDA 11.3
        suffix = '113'
    elif ver < 11050:
        # CUDA 11.4
        suffix = '114'
    elif ver < 11060:
        # CUDA 11.5
        suffix = '115'
    elif ver < 11070:
        # CUDA 11.6
        suffix = '116'
    elif ver < 11080:
        # CUDA 11.7
        suffix = '117'
    else:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too new.')
    return f'cupy-cuda{suffix}'

def _rocm_version_to_package(ver: int) -> str:
    if 400 <= ver < 410:
        # ROCm 4.0
        suffix = '4-0'
    elif 420 <= ver < 430:
        # ROCm 4.2
        suffix = '4-2'
    elif 4_03_00000 <= ver < 4_04_00000:
        # ROCm 4.3
        suffix = '4-3'
    elif 5_00_00000 <= ver < 5_01_00000:
        # ROCm 5.0
        suffix = '5-0'
    else:
        raise AutoDetectionFailed(
            f'Your ROCm version ({ver}) is unsupported.')
    return f'cupy-rocm-{suffix}'

def infer_best_package() -> str:
    """Returns the appropriate CuPy wheel package name for the environment."""

    # Find the existing CuPy wheel installation.
    installed = _find_installed_packages()
    if 1 < len(installed):
        raise AutoDetectionFailed(
            'You have multiple CuPy packages installed: \n'
            f'  {installed}\n'
            'Please uninstall all of them first, then try reinstalling.')

    elif 1 == len(installed):
        if installed[0] in PACKAGES_SDIST:
            raise AutoDetectionFailed(
                'You already have CuPy installed via source'
                ' (pip install cupy).')
        if installed[0] in PACKAGES_OUTDATED:
            raise AutoDetectionFailed(
                f'You have CuPy package "{installed[0]}" installed, but the'
                f' package is not available for version {VERSION}.')
        return installed[0]

    # Try CUDA.
    version = _get_cuda_version()
    if version is not None:
        return _cuda_version_to_package(version)

    # Try ROCm.
    version = _get_rocm_version()
    if version is not None:
        return _rocm_version_to_package(version)

    raise AutoDetectionFailed(
        'Unable to detect NVIDIA CUDA or AMD ROCm installation.')

#
# Entrypoint
#

def main() -> None:
    if os.environ.get('CUPY_UNIVERSAL_PKG_BUILD', None) is None:
        package = infer_best_package()
        requires = f'{package}=={VERSION}'
        _log(f'Installing package: {requires}')
        install_requires = [requires]
    else:
        _log('Building cupy-wheel package for release.')
        install_requires = []

    setup(
        name='cupy-wheel',
        version=VERSION,
        install_requires=install_requires,
    )

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


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2012 Roberto Alsina y otros.

# Permission is hereby granted, free of charge, to any
# person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the
# Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice
# shall be included in all copies or substantial portions of
# the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
# OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import unicode_literals, print_function

import codecs
import os
import re
import sys
import string

import unidecode
import lxml.html

from .utils import to_datetime, slugify

__all__ = ['Post']

TEASER_REGEXP = re.compile('', re.IGNORECASE)

class Post(object):

    """Represents a blog post or web page."""

    def __init__(
        self, source_path, cache_folder, destination, use_in_feeds,
        translations, default_lang, base_url, messages, template_name,
        file_metadata_regexp=None, tzinfo=None
    ):
        """Initialize post.

        The source path is the .txt post file. From it we calculate
        the meta file, as well as any translations available, and
        the .html fragment file path.
        """
        self.translated_to = set([default_lang])
        self.tags = ''
        self.date = None
        self.prev_post = None
        self.next_post = None
        self.base_url = base_url
        self.is_draft = False
        self.is_mathjax = False
        self.source_path = source_path  # posts/blah.txt
        self.post_name = os.path.splitext(source_path)[0]  # posts/blah
        # cache/posts/blah.html
        self.base_path = os.path.join(cache_folder, self.post_name + ".html")
        self.metadata_path = self.post_name + ".meta"  # posts/blah.meta
        self.folder = destination
        self.translations = translations
        self.default_lang = default_lang
        self.messages = messages
        self.template_name = template_name
        self.meta = get_meta(self.post_name, file_metadata_regexp)

        default_title = self.meta.get('title', '')
        default_pagename = self.meta.get('slug', '')
        default_description = self.meta.get('description', '')

        for k, v in self.meta.items():
            if k not in ['title', 'slug', 'description']:
                if sys.version_info[0] == 2:
                    setattr(self, unidecode.unidecode(unicode(k)), v)  # NOQA
                else:
                    setattr(self, k, v)

        if not default_title or not default_pagename or not self.date:
            raise OSError("You must set a title (found '{0}'), a slug (found "
                          "'{1}') and a date (found '{2}')! [in file "
                          "{3}]".format(default_title, default_pagename,
                                        self.date, source_path))

        # If timezone is set, build localized datetime.
        self.date = to_datetime(self.date, tzinfo)
        self.tags = [x.strip() for x in self.tags.split(',')]
        self.tags = [_f for _f in self.tags if _f]

        # While draft comes from the tags, it's not really a tag
        self.use_in_feeds = use_in_feeds and "draft" not in self.tags
        self.is_draft = 'draft' in self.tags
        self.tags = [t for t in self.tags if t != 'draft']

        # If mathjax is a tag, then enable mathjax rendering support
        self.is_mathjax = 'mathjax' in self.tags

        self.pagenames = {}
        self.titles = {}
        self.descriptions = {}

        # Load internationalized metadata
        for lang in translations:
            if lang == default_lang:
                self.titles[lang] = default_title
                self.pagenames[lang] = default_pagename
                self.descriptions[lang] = default_description
            else:
                if os.path.isfile(self.source_path + "." + lang):
                    self.translated_to.add(lang)

                meta = self.meta.copy()
                meta.update(get_meta(self.source_path, file_metadata_regexp, lang))

                # FIXME this only gets three pieces of metadata from the i18n files
                self.titles[lang] = meta['title']
                self.pagenames[lang] = meta['slug']
                self.descriptions[lang] = meta['description']

    def title(self, lang):
        """Return localized title."""
        return self.titles[lang]

    def description(self, lang):
        """Return localized description."""
        return self.descriptions[lang]

    def deps(self, lang):
        """Return a list of dependencies to build this post's page."""
        deps = [self.base_path]
        if lang != self.default_lang:
            deps += [self.base_path + "." + lang]
        deps += self.fragment_deps(lang)
        return deps

    def fragment_deps(self, lang):
        """Return a list of dependencies to build this post's fragment."""
        deps = [self.source_path]
        if os.path.isfile(self.metadata_path):
            deps.append(self.metadata_path)
        if lang != self.default_lang:
            lang_deps = list(filter(os.path.exists, [x + "." + lang for x in
                                                     deps]))
            deps += lang_deps
        return deps

    def is_translation_available(self, lang):
        """Return true if the translation actually exists."""
        return lang in self.translated_to

    def _translated_file_path(self, lang):
        """Return path to the translation's file, or to the original."""
        file_name = self.base_path
        if lang != self.default_lang:
            file_name_lang = '.'.join((file_name, lang))
            if os.path.exists(file_name_lang):
                file_name = file_name_lang
        return file_name

    def text(self, lang, teaser_only=False, strip_html=False):
        """Read the post file for that language and return its contents"""
        file_name = self._translated_file_path(lang)

        with codecs.open(file_name, "r", "utf8") as post_file:
            data = post_file.read()

        if data:
            data = lxml.html.make_links_absolute(data, self.permalink(lang=lang))
        if data and teaser_only:
            e = lxml.html.fromstring(data)
            teaser = []
            teaser_str = self.messages[lang]["Read more"] + '...'
            flag = False
            for elem in e:
                elem_string = lxml.html.tostring(elem).decode('utf8')
                match = TEASER_REGEXP.match(elem_string)
                if match:
                    flag = True
                    if match.group(2):
                        teaser_str = match.group(2)
                    break
                teaser.append(elem_string)
            if flag:
                teaser.append('
{1}
'.format(
                    self.permalink(lang), teaser_str))
            data = ''.join(teaser)

        if data and strip_html:
            content = lxml.html.fromstring(data)
            data = content.text_content().strip()  # No whitespace wanted.

        return data

    def destination_path(self, lang, extension='.html'):
        path = os.path.join(self.translations[lang],
                            self.folder, self.pagenames[lang] + extension)
        return path

    def permalink(self, lang=None, absolute=False, extension='.html'):
        if lang is None:
            lang = self.default_lang
        pieces = list(os.path.split(self.translations[lang]))
        pieces += list(os.path.split(self.folder))
        pieces += [self.pagenames[lang] + extension]
        pieces = [_f for _f in pieces if _f and _f != '.']
        if absolute:
            pieces = [self.base_url] + pieces
        else:
            pieces = [""] + pieces
        link = "/".join(pieces)
        return link

    def source_ext(self):
        return os.path.splitext(self.source_path)[1]

# Code that fetches metadata from different places

def re_meta(line, match=None):
    """re.compile for meta"""
    if match:
        reStr = re.compile('^\.\. {0}: (.*)'.format(re.escape(match)))
    else:
        reStr = re.compile('^\.\. (.*?): (.*)')
    result = reStr.findall(line.strip())
    if match and result:
        return (match, result[0])
    elif not match and result:
        return (result[0][0], result[0][1].strip())
    else:
        return (None,)

def _get_metadata_from_filename_by_regex(filename, metadata_regexp):
    """
    Tries to ried the metadata from the filename based on the given re.
    This requires to use symbolic group names in the pattern.

    The part to read the metadata from the filename based on a regular
    expression is taken from Pelican - pelican/readers.py
    """
    match = re.match(metadata_regexp, filename)
    meta = {}

    if match:
        # .items() for py3k compat.
        for key, value in match.groupdict().items():
            meta[key.lower()] = value  # metadata must be lowercase

    return meta

def get_metadata_from_file(source_path):
    """Extracts metadata from the file itself, by parsing contents."""
    try:
        with codecs.open(source_path, "r", "utf8") as meta_file:
            meta_data = [x.strip() for x in meta_file.readlines()]
        return _get_metadata_from_file(meta_data)
    except Exception:  # The file may not exist, for multilingual sites
        return {}

def _get_metadata_from_file(meta_data):
    """Parse file contents and obtain metadata.

    >>> g = _get_metadata_from_file
    >>> list(g([]).values())
    []
    >>> str(g(["FooBar","======"])["title"])
    'FooBar'
    >>> str(g(["#FooBar"])["title"])
    'FooBar'
    >>> str(g([".. title: FooBar"])["title"])
    'FooBar'
    >>> 'title' in g(["",".. title: FooBar"])
    False

    """
    meta = {}

    re_md_title = re.compile(r'^{0}([^{0}].*)'.format(re.escape('#')))
    # Assuming rst titles are going to be at least 4 chars long
    # otherwise this detects things like ''' wich breaks other markups.
    re_rst_title = re.compile(r'^([{0}]{{4,}})'.format(re.escape(
        string.punctuation)))

    for i, line in enumerate(meta_data):
        if not line:
            break
        if 'title' not in meta:
            match = re_meta(line, 'title')
            if match[0]:
                meta['title'] = match[1]
        if 'title' not in meta:
            if re_rst_title.findall(line) and i > 0:
                meta['title'] = meta_data[i - 1].strip()
        if 'title' not in meta:
            if re_md_title.findall(line):
                meta['title'] = re_md_title.findall(line)[0]

        match = re_meta(line)
        if match[0]:
            meta[match[0]] = match[1]

    return meta

def get_metadata_from_meta_file(path, lang=None):
    """Takes a post path, and gets data from a matching .meta file."""

    meta_path = os.path.splitext(path)[0] + '.meta'
    if lang:
        meta_path += '.' + lang
    if os.path.isfile(meta_path):
        with codecs.open(meta_path, "r", "utf8") as meta_file:
            meta_data = meta_file.readlines()
        while len(meta_data) < 6:
            meta_data.append("")
        (title, slug, date, tags, link, description) = [
            x.strip() for x in meta_data][:6]

        meta = {}

        if title: meta['title'] = title
        if slug: meta['slug'] = slug
        if date: meta['date'] = date
        if tags: meta['tags'] = tags
        if link: meta['link'] = link
        if description: meta['description'] = description

        return meta
    else:
        return {}

def get_meta(source_path, file_metadata_regexp=None, lang=None):
    """Get post's meta from source.

    If ``file_metadata_regexp`` is given it will be tried to read
    metadata from the filename.
    If any metadata is then found inside the file the metadata from the
    file will override previous findings.
    """
    meta = {}

    meta.update(get_metadata_from_meta_file(source_path, lang))

    if meta:
        return meta

    if file_metadata_regexp is not None:
        meta.update(_get_metadata_from_filename_by_regex(source_path,
                                                         file_metadata_regexp))

    meta.update(get_metadata_from_file(source_path))

    if 'slug' not in meta:
        # If no slug is found in the metadata use the filename
        meta['slug'] = slugify(os.path.splitext(
            os.path.basename(source_path))[0])

    if 'title' not in meta:
        # If no title is found, use the filename without extension
        meta['title'] = os.path.splitext(os.path.basename(source_path))[0]

    return meta
<commit_msg>read metadata from the right place (also makes get_meta almost a method of post ;-)<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2012 Roberto Alsina y otros.

# Permission is hereby granted, free of charge, to any
# person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the
# Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice
# shall be included in all copies or substantial portions of
# the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
# OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import unicode_literals, print_function

import codecs
import os
import re
import sys
import string

import unidecode
import lxml.html

from .utils import to_datetime, slugify

__all__ = ['Post']

TEASER_REGEXP = re.compile('', re.IGNORECASE)

class Post(object):

    """Represents a blog post or web page."""

    def __init__(
        self, source_path, cache_folder, destination, use_in_feeds,
        translations, default_lang, base_url, messages, template_name,
        file_metadata_regexp=None, tzinfo=None
    ):
        """Initialize post.

        The source path is the .txt post file. From it we calculate
        the meta file, as well as any translations available, and
        the .html fragment file path.
        """
        self.translated_to = set([default_lang])
        self.tags = ''
        self.date = None
        self.prev_post = None
        self.next_post = None
        self.base_url = base_url
        self.is_draft = False
        self.is_mathjax = False
        self.source_path = source_path  # posts/blah.txt
        self.post_name = os.path.splitext(source_path)[0]  # posts/blah
        # cache/posts/blah.html
        self.base_path = os.path.join(cache_folder, self.post_name + ".html")
        self.metadata_path = self.post_name + ".meta"  # posts/blah.meta
        self.folder = destination
        self.translations = translations
        self.default_lang = default_lang
        self.messages = messages
        self.template_name = template_name
        self.meta = get_meta(self, file_metadata_regexp)

        default_title = self.meta.get('title', '')
        default_pagename = self.meta.get('slug', '')
        default_description = self.meta.get('description', '')

        for k, v in self.meta.items():
            if k not in ['title', 'slug', 'description']:
                if sys.version_info[0] == 2:
                    setattr(self, unidecode.unidecode(unicode(k)), v)  # NOQA
                else:
                    setattr(self, k, v)

        if not default_title or not default_pagename or not self.date:
            raise OSError("You must set a title (found '{0}'), a slug (found "
                          "'{1}') and a date (found '{2}')! [in file "
                          "{3}]".format(default_title, default_pagename,
                                        self.date, source_path))

        # If timezone is set, build localized datetime.
        self.date = to_datetime(self.date, tzinfo)
        self.tags = [x.strip() for x in self.tags.split(',')]
        self.tags = [_f for _f in self.tags if _f]

        # While draft comes from the tags, it's not really a tag
        self.use_in_feeds = use_in_feeds and "draft" not in self.tags
        self.is_draft = 'draft' in self.tags
        self.tags = [t for t in self.tags if t != 'draft']

        # If mathjax is a tag, then enable mathjax rendering support
        self.is_mathjax = 'mathjax' in self.tags

        self.pagenames = {}
        self.titles = {}
        self.descriptions = {}

        # Load internationalized metadata
        for lang in translations:
            if lang == default_lang:
                self.titles[lang] = default_title
                self.pagenames[lang] = default_pagename
                self.descriptions[lang] = default_description
            else:
                if os.path.isfile(self.source_path + "." + lang):
                    self.translated_to.add(lang)

                meta = self.meta.copy()
                meta.update(get_meta(self.source_path, file_metadata_regexp, lang))

                # FIXME this only gets three pieces of metadata from the i18n files
                self.titles[lang] = meta['title']
                self.pagenames[lang] = meta['slug']
                self.descriptions[lang] = meta['description']

    def title(self, lang):
        """Return localized title."""
        return self.titles[lang]

    def description(self, lang):
        """Return localized description."""
        return self.descriptions[lang]

    def deps(self, lang):
        """Return a list of dependencies to build this post's page."""
        deps = [self.base_path]
        if lang != self.default_lang:
            deps += [self.base_path + "." + lang]
        deps += self.fragment_deps(lang)
        return deps

    def fragment_deps(self, lang):
        """Return a list of dependencies to build this post's fragment."""
        deps = [self.source_path]
        if os.path.isfile(self.metadata_path):
            deps.append(self.metadata_path)
        if lang != self.default_lang:
            lang_deps = list(filter(os.path.exists, [x + "." + lang for x in
                                                     deps]))
            deps += lang_deps
        return deps

    def is_translation_available(self, lang):
        """Return true if the translation actually exists."""
        return lang in self.translated_to

    def _translated_file_path(self, lang):
        """Return path to the translation's file, or to the original."""
        file_name = self.base_path
        if lang != self.default_lang:
            file_name_lang = '.'.join((file_name, lang))
            if os.path.exists(file_name_lang):
                file_name = file_name_lang
        return file_name

    def text(self, lang, teaser_only=False, strip_html=False):
        """Read the post file for that language and return its contents"""
        file_name = self._translated_file_path(lang)

        with codecs.open(file_name, "r", "utf8") as post_file:
            data = post_file.read()

        if data:
            data = lxml.html.make_links_absolute(data, self.permalink(lang=lang))
        if data and teaser_only:
            e = lxml.html.fromstring(data)
            teaser = []
            teaser_str = self.messages[lang]["Read more"] + '...'
            flag = False
            for elem in e:
                elem_string = lxml.html.tostring(elem).decode('utf8')
                match = TEASER_REGEXP.match(elem_string)
                if match:
                    flag = True
                    if match.group(2):
                        teaser_str = match.group(2)
                    break
                teaser.append(elem_string)
            if flag:
                teaser.append('
{1}
'.format(
                    self.permalink(lang), teaser_str))
            data = ''.join(teaser)

        if data and strip_html:
            content = lxml.html.fromstring(data)
            data = content.text_content().strip()  # No whitespace wanted.

        return data

    def destination_path(self, lang, extension='.html'):
        path = os.path.join(self.translations[lang],
                            self.folder, self.pagenames[lang] + extension)
        return path

    def permalink(self, lang=None, absolute=False, extension='.html'):
        if lang is None:
            lang = self.default_lang
        pieces = list(os.path.split(self.translations[lang]))
        pieces += list(os.path.split(self.folder))
        pieces += [self.pagenames[lang] + extension]
        pieces = [_f for _f in pieces if _f and _f != '.']
        if absolute:
            pieces = [self.base_url] + pieces
        else:
            pieces = [""] + pieces
        link = "/".join(pieces)
        return link

    def source_ext(self):
        return os.path.splitext(self.source_path)[1]

# Code that fetches metadata from different places

def re_meta(line, match=None):
    """re.compile for meta"""
    if match:
        reStr = re.compile('^\.\. {0}: (.*)'.format(re.escape(match)))
    else:
        reStr = re.compile('^\.\. (.*?): (.*)')
    result = reStr.findall(line.strip())
    if match and result:
        return (match, result[0])
    elif not match and result:
        return (result[0][0], result[0][1].strip())
    else:
        return (None,)

def _get_metadata_from_filename_by_regex(filename, metadata_regexp):
    """
    Tries to ried the metadata from the filename based on the given re.
    This requires to use symbolic group names in the pattern.

    The part to read the metadata from the filename based on a regular
    expression is taken from Pelican - pelican/readers.py
    """
    match = re.match(metadata_regexp, filename)
    meta = {}

    if match:
        # .items() for py3k compat.
        for key, value in match.groupdict().items():
            meta[key.lower()] = value  # metadata must be lowercase

    return meta

def get_metadata_from_file(source_path):
    """Extracts metadata from the file itself, by parsing contents."""
    try:
        with codecs.open(source_path, "r", "utf8") as meta_file:
            meta_data = [x.strip() for x in meta_file.readlines()]
        return _get_metadata_from_file(meta_data)
    except Exception:  # The file may not exist, for multilingual sites
        return {}

def _get_metadata_from_file(meta_data):
    """Parse file contents and obtain metadata.

    >>> g = _get_metadata_from_file
    >>> list(g([]).values())
    []
    >>> str(g(["FooBar","======"])["title"])
    'FooBar'
    >>> str(g(["#FooBar"])["title"])
    'FooBar'
    >>> str(g([".. title: FooBar"])["title"])
    'FooBar'
    >>> 'title' in g(["",".. title: FooBar"])
    False

    """
    meta = {}

    re_md_title = re.compile(r'^{0}([^{0}].*)'.format(re.escape('#')))
    # Assuming rst titles are going to be at least 4 chars long
    # otherwise this detects things like ''' wich breaks other markups.
    re_rst_title = re.compile(r'^([{0}]{{4,}})'.format(re.escape(
        string.punctuation)))

    for i, line in enumerate(meta_data):
        if not line:
            break
        if 'title' not in meta:
            match = re_meta(line, 'title')
            if match[0]:
                meta['title'] = match[1]
        if 'title' not in meta:
            if re_rst_title.findall(line) and i > 0:
                meta['title'] = meta_data[i - 1].strip()
        if 'title' not in meta:
            if re_md_title.findall(line):
                meta['title'] = re_md_title.findall(line)[0]

        match = re_meta(line)
        if match[0]:
            meta[match[0]] = match[1]

    return meta

def get_metadata_from_meta_file(path, lang=None):
    """Takes a post path, and gets data from a matching .meta file."""
    meta_path = os.path.splitext(path)[0] + '.meta'
    if lang:
        meta_path += '.' + lang
    if os.path.isfile(meta_path):
        with codecs.open(meta_path, "r", "utf8") as meta_file:
            meta_data = meta_file.readlines()
        while len(meta_data) < 6:
            meta_data.append("")
        (title, slug, date, tags, link, description) = [
            x.strip() for x in meta_data][:6]

        meta = {}

        if title: meta['title'] = title
        if slug: meta['slug'] = slug
        if date: meta['date'] = date
        if tags: meta['tags'] = tags
        if link: meta['link'] = link
        if description: meta['description'] = description

        return meta
    else:
        return {}

def get_meta(post, file_metadata_regexp=None, lang=None):
    """Get post's meta from source.

    If ``file_metadata_regexp`` is given it will be tried to read
    metadata from the filename.
    If any metadata is then found inside the file the metadata from the
    file will override previous findings.
    """
    meta = {}

    meta.update(get_metadata_from_meta_file(post.metadata_path, lang))

    if meta:
        return meta

    if file_metadata_regexp is not None:
        meta.update(_get_metadata_from_filename_by_regex(post.source_path,
                                                         file_metadata_regexp))

    meta.update(get_metadata_from_file(post.source_path))

    if 'slug' not in meta:
        # If no slug is found in the metadata use the filename
        meta['slug'] = slugify(os.path.splitext(
            os.path.basename(source_path))[0])

    if 'title' not in meta:
        # If no title is found, use the filename without extension
        meta['title'] = os.path.splitext(os.path.basename(source_path))[0]

    return meta
<|endoftext|>


<commit_before>"""Defines the class that represents running job executions"""
from __future__ import unicode_literals

import logging
import threading

from django.db import transaction
from django.utils.timezone import now

from error.models import Error
from job.execution.tasks.main_task import MainTask
from job.execution.tasks.post_task import PostTask
from job.execution.tasks.pre_task import PreTask
from job.execution.tasks.pull_task import PullTask
from job.models import JobExecution, JobExecutionEnd
from job.tasks.update import TaskStatusUpdate
from util.retry import retry_database_query

logger = logging.getLogger(__name__)

class RunningJobExecution(object):
    """This class represents a currently running job execution. This class is thread-safe."""

    def __init__(self, agent_id, job_exe, job_type, configuration):
        """Constructor

        :param agent_id: The ID of the agent on which the execution is running
        :type agent_id: string
        :param job_exe: The job execution model, related fields will only have IDs populated
        :type job_exe: :class:`job.models.JobExecution`
        :param job_type: The job type model
        :type job_type: :class:`job.models.JobType`
        :param configuration: The job execution configuration, including secret values
        :type configuration: :class:`job.configuration.json.execution.exe_config.ExecutionConfiguration`
        """

        # Public, read-only info
        self.id = job_exe.id
        self.cluster_id = job_exe.get_cluster_id()
        self.job_id = job_exe.job_id
        self.exe_num = job_exe.exe_num
        self.job_type_id = job_exe.job_type_id
        self.node_id = job_exe.node_id
        self.queued = job_exe.queued
        self.started = job_exe.started
        self.docker_volumes = configuration.get_named_docker_volumes()

        # Internal task and status info
        self._lock = threading.Lock()  # Protects the following fields
        self._all_tasks = []
        self._current_task = None
        self._error = None
        self._finished = None
        self._remaining_tasks = []
        self._status = 'RUNNING'

        # Create tasks
        for task_type in configuration.get_task_types():
            task = None
            if task_type == 'pull':
                task = PullTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'pre':
                task = PreTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'main':
                task = MainTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'post':
                task = PostTask(agent_id, job_exe, job_type, configuration)
            self._all_tasks.append(task)
        for task in self._all_tasks:
            self._remaining_tasks.append(task)

    @property
    def current_task(self):
        """Returns the currently running task of the job execution, or None if no task is currently running

        :returns: The current task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        return self._current_task

    @property
    def error_category(self):
        """Returns the category of this job execution's error, None if there is no error

        :returns: The error category, possibly None
        :rtype: bool
        """

        if self._error:
            return self._error.category
        return None

    @property
    def finished(self):
        """When this job execution finished, possibly None

        :returns: When this job execution finished, possibly None
        :rtype: :class:`datetime.datetime`
        """

        return self._finished

    @property
    def status(self):
        """Returns the status of this job execution

        :returns: The status of the job execution
        :rtype: string
        """

        return self._status

    def create_job_exe_end_model(self):
        """Creates and returns a job execution end model for this job execution. Caller must ensure that this job
        execution is finished before calling.

        :returns: The job execution end model
        :rtype: :class:`job.models.JobExecutionEnd`
        """

        job_exe_end = JobExecutionEnd()
        job_exe_end.job_exe_id = self.id
        job_exe_end.job_id = self.job_id
        job_exe_end.job_type = self.job_type_id
        job_exe_end.exe_num = self.exe_num
        job_exe_end.status = self._status
        if self._error:
            job_exe_end.error_id = self._error.id
        job_exe_end.node_id = self.node_id
        job_exe_end.queued = self.queued
        job_exe_end.started = self.started
        job_exe_end.ended = self._finished
        return job_exe_end

    def execution_canceled(self):
        """Cancels this job execution and returns the current task

        :returns: The current task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            task = self._current_task
            self._finish_execution('CANCELED', now())
            return task

    @retry_database_query
    def execution_lost(self, when):
        """Fails this job execution for its node becoming lost and returns the current task

        :param when: The time that the node was lost
        :type when: :class:`datetime.datetime`
        """

        error = Error.objects.get_builtin_error('node-lost')
        from queue.models import Queue
        Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._current_task = None
            self._remaining_tasks = []
            self._set_finished_status('FAILED', when, error)

    @retry_database_query
    def execution_timed_out(self, task, when):
        """Fails this job execution for timing out

        :param task: The task that timed out
        :type task: :class:`job.tasks.exe_task.JobExecutionTask`
        :param when: The time that the job execution timed out
        :type when: :class:`datetime.datetime`
        """

        if task.has_started:
            error_name = task.timeout_error_name
        else:
            error_name = 'launch-timeout'
        error = Error.objects.get_builtin_error(error_name)
        from queue.models import Queue
        Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._current_task = None
            self._remaining_tasks = []
            self._set_finished_status('FAILED', when, error)

    def get_container_names(self):
        """Returns the list of container names for all tasks in this job execution

        :returns: The list of all container names
        :rtype: [string]
        """

        containers = []
        with self._lock:
            for task in self._all_tasks:
                if task.container_name:
                    containers.append(task.container_name)
            return containers

    def is_finished(self):
        """Indicates whether this job execution is finished with all tasks

        :returns: True if all tasks are finished, False otherwise
        :rtype: bool
        """

        with self._lock:
            return not self._current_task and not self._remaining_tasks

    def is_next_task_ready(self):
        """Indicates whether the next task in this job execution is ready

        :returns: True if the next task is ready, False otherwise
        :rtype: bool
        """

        with self._lock:
            return not self._current_task and self._remaining_tasks

    def next_task(self):
        """Returns the next task in this job execution. Returns None if there are no remaining tasks.

        :returns: The next task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            if not self._remaining_tasks:
                return None

            next_task = self._remaining_tasks[0]
            return next_task

    def start_next_task(self):
        """Starts the next task in the job execution and returns it. Returns None if the next task is not ready or no
        tasks remain.

        :returns: The new task that was started, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            if self._current_task or not self._remaining_tasks:
                return None

            self._current_task = self._remaining_tasks.pop(0)
            return self._current_task

    def task_update(self, task_update):
        """Updates a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        if task_update.status == TaskStatusUpdate.FINISHED:
            self._task_complete(task_update)
        elif task_update.status == TaskStatusUpdate.LOST:
            self._task_lost(task_update)
        elif task_update.status in [TaskStatusUpdate.FAILED, TaskStatusUpdate.KILLED]:
            self._task_fail(task_update)

    def _finish_execution(self, status, when, error=None):
        """Sets the finished status for this job execution. Caller must have obtained lock.

        :param status: The status
        :type status: string
        :param when: The time that the job execution finished
        :type when: :class:`datetime.datetime`
        :param error: The error, possibly None
        :type error: :class:`error.models.Error`
        """

        if self._status == 'RUNNING':
            self._current_task = None
            self._remaining_tasks = []
            self._status = status
            self._finished = when
            self._error = error

    @retry_database_query
    def _task_complete(self, task_update):
        """Completes a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        with self._lock:
            current_task = self._current_task
            remaining_tasks = self._remaining_tasks

        if not current_task or current_task.id != task_update.task_id:
            return

        when = now()
        with transaction.atomic():
            need_refresh = current_task.complete(task_update)
            if need_refresh and remaining_tasks:
                job_exe = JobExecution.objects.get(id=self.id)
                for task in remaining_tasks:
                    task.refresh_cached_values(job_exe)
            if not remaining_tasks:
                from queue.models import Queue
                Queue.objects.handle_job_completion(self.id, when, self._all_tasks)

        with self._lock:
            if self._current_task and self._current_task.id == task_update.task_id:
                self._current_task = None
                if not self._remaining_tasks:
                    self._set_finished_status('COMPLETED', when)

    @retry_database_query
    def _task_fail(self, task_update):
        """Fails a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        current_task = self._current_task
        if not current_task or current_task.id != task_update.task_id:
            return

        when = now()
        with transaction.atomic():
            error = current_task.determine_error(task_update)
            from queue.models import Queue
            Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._current_task = None
            self._remaining_tasks = []
            self._set_finished_status('FAILED', when, error)

    def _task_lost(self, task_update):
        """Tells this job execution that one of its tasks was lost

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        with self._lock:
            if not self._current_task or self._current_task.id != task_update.task_id:
                return

            self._current_task.update_task_id_for_lost_task()  # Note: This changes the task ID!
            self._remaining_tasks.insert(0, self._current_task)
            self._current_task = None
<commit_msg>Minor fixes<commit_after>"""Defines the class that represents running job executions"""
from __future__ import unicode_literals

import logging
import threading

from django.db import transaction
from django.utils.timezone import now

from error.models import Error
from job.execution.tasks.main_task import MainTask
from job.execution.tasks.post_task import PostTask
from job.execution.tasks.pre_task import PreTask
from job.execution.tasks.pull_task import PullTask
from job.models import JobExecution, JobExecutionEnd
from job.tasks.update import TaskStatusUpdate
from util.retry import retry_database_query

logger = logging.getLogger(__name__)

class RunningJobExecution(object):
    """This class represents a currently running job execution. This class is thread-safe."""

    def __init__(self, agent_id, job_exe, job_type, configuration):
        """Constructor

        :param agent_id: The ID of the agent on which the execution is running
        :type agent_id: string
        :param job_exe: The job execution model, related fields will only have IDs populated
        :type job_exe: :class:`job.models.JobExecution`
        :param job_type: The job type model
        :type job_type: :class:`job.models.JobType`
        :param configuration: The job execution configuration, including secret values
        :type configuration: :class:`job.configuration.json.execution.exe_config.ExecutionConfiguration`
        """

        # Public, read-only info
        self.id = job_exe.id
        self.cluster_id = job_exe.get_cluster_id()
        self.job_id = job_exe.job_id
        self.exe_num = job_exe.exe_num
        self.job_type_id = job_exe.job_type_id
        self.node_id = job_exe.node_id
        self.queued = job_exe.queued
        self.started = job_exe.started
        self.docker_volumes = configuration.get_named_docker_volumes()

        # Internal task and status info
        self._lock = threading.Lock()  # Protects the following fields
        self._all_tasks = []
        self._current_task = None
        self._error = None
        self._finished = None
        self._remaining_tasks = []
        self._status = 'RUNNING'

        # Create tasks
        for task_type in configuration.get_task_types():
            task = None
            if task_type == 'pull':
                task = PullTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'pre':
                task = PreTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'main':
                task = MainTask(agent_id, job_exe, job_type, configuration)
            elif task_type == 'post':
                task = PostTask(agent_id, job_exe, job_type, configuration)
            self._all_tasks.append(task)
        for task in self._all_tasks:
            self._remaining_tasks.append(task)

    @property
    def current_task(self):
        """Returns the currently running task of the job execution, or None if no task is currently running

        :returns: The current task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        return self._current_task

    @property
    def error_category(self):
        """Returns the category of this job execution's error, None if there is no error

        :returns: The error category, possibly None
        :rtype: bool
        """

        if self._error:
            return self._error.category
        return None

    @property
    def finished(self):
        """When this job execution finished, possibly None

        :returns: When this job execution finished, possibly None
        :rtype: :class:`datetime.datetime`
        """

        return self._finished

    @property
    def status(self):
        """Returns the status of this job execution

        :returns: The status of the job execution
        :rtype: string
        """

        return self._status

    def create_job_exe_end_model(self):
        """Creates and returns a job execution end model for this job execution. Caller must ensure that this job
        execution is finished before calling.

        :returns: The job execution end model
        :rtype: :class:`job.models.JobExecutionEnd`
        """

        job_exe_end = JobExecutionEnd()
        job_exe_end.job_exe_id = self.id
        job_exe_end.job_id = self.job_id
        job_exe_end.job_type = self.job_type_id
        job_exe_end.exe_num = self.exe_num
        job_exe_end.status = self._status
        if self._error:
            job_exe_end.error_id = self._error.id
        job_exe_end.node_id = self.node_id
        job_exe_end.queued = self.queued
        job_exe_end.started = self.started
        job_exe_end.ended = self._finished
        return job_exe_end

    def execution_canceled(self):
        """Cancels this job execution and returns the current task

        :returns: The current task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            task = self._current_task
            self._set_final_status('CANCELED', now())
            return task

    @retry_database_query
    def execution_lost(self, when):
        """Fails this job execution for its node becoming lost

        :param when: The time that the node was lost
        :type when: :class:`datetime.datetime`
        """

        error = Error.objects.get_builtin_error('node-lost')
        from queue.models import Queue
        Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._current_task = None
            self._set_final_status('FAILED', when, error)

    @retry_database_query
    def execution_timed_out(self, task, when):
        """Fails this job execution for timing out

        :param task: The task that timed out
        :type task: :class:`job.tasks.exe_task.JobExecutionTask`
        :param when: The time that the job execution timed out
        :type when: :class:`datetime.datetime`
        """

        if task.has_started:
            error_name = task.timeout_error_name
        else:
            error_name = 'launch-timeout'
        error = Error.objects.get_builtin_error(error_name)
        from queue.models import Queue
        Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._set_final_status('FAILED', when, error)

    def get_container_names(self):
        """Returns the list of container names for all tasks in this job execution

        :returns: The list of all container names
        :rtype: [string]
        """

        containers = []
        with self._lock:
            for task in self._all_tasks:
                if task.container_name:
                    containers.append(task.container_name)
            return containers

    def is_finished(self):
        """Indicates whether this job execution is finished with all tasks

        :returns: True if all tasks are finished, False otherwise
        :rtype: bool
        """

        with self._lock:
            return not self._current_task and not self._remaining_tasks

    def is_next_task_ready(self):
        """Indicates whether the next task in this job execution is ready

        :returns: True if the next task is ready, False otherwise
        :rtype: bool
        """

        with self._lock:
            return not self._current_task and self._remaining_tasks

    def next_task(self):
        """Returns the next task in this job execution. Returns None if there are no remaining tasks.

        :returns: The next task, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            if not self._remaining_tasks:
                return None

            next_task = self._remaining_tasks[0]
            return next_task

    def start_next_task(self):
        """Starts the next task in the job execution and returns it. Returns None if the next task is not ready or no
        tasks remain.

        :returns: The new task that was started, possibly None
        :rtype: :class:`job.tasks.base_task.Task`
        """

        with self._lock:
            if self._current_task or not self._remaining_tasks:
                return None

            self._current_task = self._remaining_tasks.pop(0)
            return self._current_task

    def task_update(self, task_update):
        """Updates a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        if task_update.status == TaskStatusUpdate.FINISHED:
            self._task_complete(task_update)
        elif task_update.status == TaskStatusUpdate.LOST:
            self._task_lost(task_update)
        elif task_update.status in [TaskStatusUpdate.FAILED, TaskStatusUpdate.KILLED]:
            self._task_fail(task_update)

    def _set_final_status(self, status, when, error=None):
        """Sets the final status for this job execution and removes all remaining tasks. The current task remains since
        it may need to be killed. Caller must have obtained lock.

        :param status: The status
        :type status: string
        :param when: The time that the job execution finished
        :type when: :class:`datetime.datetime`
        :param error: The error, possibly None
        :type error: :class:`error.models.Error`
        """

        if self._status == 'RUNNING':
            self._remaining_tasks = []
            self._status = status
            self._finished = when
            self._error = error

    @retry_database_query
    def _task_complete(self, task_update):
        """Completes a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        with self._lock:
            current_task = self._current_task
            remaining_tasks = self._remaining_tasks

        if not current_task or current_task.id != task_update.task_id:
            return

        when = now()
        with transaction.atomic():
            need_refresh = current_task.complete(task_update)
            if need_refresh and remaining_tasks:
                job_exe = JobExecution.objects.get(id=self.id)
                for task in remaining_tasks:
                    task.refresh_cached_values(job_exe)
            if not remaining_tasks:
                from queue.models import Queue
                Queue.objects.handle_job_completion(self.id, when, self._all_tasks)

        with self._lock:
            if self._current_task and self._current_task.id == task_update.task_id:
                self._current_task = None
                if not self._remaining_tasks:
                    self._set_finished_status('COMPLETED', when)

    @retry_database_query
    def _task_fail(self, task_update):
        """Fails a task for this job execution

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        current_task = self._current_task
        if not current_task or current_task.id != task_update.task_id:
            return

        when = now()
        with transaction.atomic():
            error = current_task.determine_error(task_update)
            from queue.models import Queue
            Queue.objects.handle_job_failure(self.id, when, self._all_tasks, error)

        with self._lock:
            self._current_task = None
            self._remaining_tasks = []
            self._set_finished_status('FAILED', when, error)

    def _task_lost(self, task_update):
        """Tells this job execution that one of its tasks was lost

        :param task_update: The task update
        :type task_update: :class:`job.tasks.update.TaskStatusUpdate`
        """

        with self._lock:
            if not self._current_task or self._current_task.id != task_update.task_id:
                return

            self._current_task.update_task_id_for_lost_task()  # Note: This changes the task ID!
            self._remaining_tasks.insert(0, self._current_task)
            self._current_task = None
<|endoftext|>


<commit_before>"""Handles secret getters and setters for Scale"""

import ast
import jwt
import json
import requests

from django.conf import settings
from vault.exceptions import InvalidSecretsAuthorization, InvalidSecretsRequest, InvalidSecretsToken, InvalidSecretsValue

class SecretsHandler(object):
    """Represents a secrets handler for setting and retrieving secrets
    """

    def __init__(self):
        """Creates a secrets handler object.  The backend is initially tested to ensure it exists and Scale can
        authenticate properly with it.
        """

        self.secrets_error_codes = {
            301: 'Status Code 301 - Too many "/" separators in URL.',
            400: 'Status Code 400 - Invalid request, missing or invalid data.',
            403: 'Status Code 403 - Forbidden, authentication details are incorrect or no access to feature.',
            404: 'Status Code 404 - Invalid path, secret or backend not found.',
            405: 'Status Code 405 - Unsupported operation.',
            409: 'Status Code 409 - Secret or secret store already exists.',
            429: 'Status Code 429 - Rate limit exceeded.',
            500: 'Status Code 500 - Internal server error.',
            503: 'Status Code 503 - Secrets store is down for maintenance or is currently sealed.',
        }

        self.secrets_url = settings.SECRETS_URL
        self.secrets_token = settings.SECRETS_TOKEN
        self.service_account = settings.DCOS_SERVICE_ACCOUNT

        if self.service_account:
            self.dcos_token = self._dcos_authenticate()
        else:
            self.dcos_token = None
            self._vault_authenticate()

        self._check_secrets_backend()

    def get_job_type_secrets(self, job_name):
        """Retrieves the secrets located at the job_name within the backend

        :param job_name: path within the secrets backend that the secret is stored
        :type job_name: str
        :return: secret_values
        :rtype: str
        """

        url = self.secrets_url

        if self.dcos_token:
            url = ''.join([url, 'secret/default/scale/job-type/', job_name])
            data = json.dumps({'uid': self.service_account, 'token': self.dcos_token})
            get_secret = self._make_request('GET', url, data=data)

            try:
                response = get_secret.json()
                secret_values = ast.literal_eval(response['value'])
            except SyntaxError:
                raise InvalidSecretsValue('The DCOS secrets value could not be converted to JSON.')

        else:
            url = url + 'secret/scale/job-type/' + job_name
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            get_secret = self._make_request('GET', url, headers=headers)

            response = get_secret.json()
            secret_values = response['data']

        return secret_values

    def list_job_types(self):
        """Gets the names of all job types that have secrets

        :return: all_job_types
        :rtype: list
        """

        url = self.secrets_url

        if self.dcos_token:
            url += 'secret/default/scale/job-type?list=true'
            data = json.dumps({'uid': self.service_account, 'token': self.dcos_token})
            list_jobs = self._make_request('GET', url, data=data)

            response = list_jobs.json()
            all_job_types = response['array']

        else:
            url += 'secret/scale/job-type'
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            list_jobs = self._make_request('LIST', url, headers)

            response = list_jobs.json()
            all_job_types = response['data']['keys']

        return all_job_types

    def set_job_type_secrets(self, secret_path, secret_json):
        """write job-type secrets to the secrets backend

        :param secret_path: path within the secrets backend that the secret will be stored.
        :type secret_path: str
        :param secret_json: dict with name:value pairs for all secrets associated with the job
        :type secret_json: dict
        :return:
        """

        url = self.secrets_url
        secret_values = json.dumps(secret_json)

        if self.dcos_token:
            url = ''.join([url, 'secret/default/scale/job-type', secret_path])
            data = json.dumps({
                'uid': self.service_account,
                'token': self.dcos_token,
                'author': 'scale',
                'value': secret_values
            })
            set_secret = self._make_request('PUT', url, data=data)

        else:
            url = ''.join([url, 'secret/scale/job-type', secret_path])
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            data = secret_values
            set_secret = self._make_request('PUT', url, headers, data)

    def _check_secrets_backend(self):
        """Validates that Scale can transact with the secrets backend properly.
        """

        url = self.secrets_url

        if self.dcos_token:
            pass

        else:
            url += '/sys/mounts'
            headers = {
                "Content-Type": "application/json",
                "X-Vault-Token": self.secrets_token
            }

            check_mount = self._make_request('GET', url, headers=headers, catch_error=False)

            if check_mount.status_code == 200:
                if 'scale/' not in json.loads(check_mount.content).keys():
                    self._create_vault_store()
            elif check_mount.status_code == 403:
                raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
            elif check_mount.status_code == 404:
                self._create_dcos_store()

    # def _create_dcos_store(self):
    #     """Create a new store within the DC/OS Vault backend that will be used for all Scale secrets.

    #     **As of 9FEB2017 DC/OS does not support additional backends.**
    #     """

    #     url = self.secrets_url + '/store/scale'
    #     data = json.dumps({
    #         'uid': self.service_account,
    #         'token': self.dcos_token,
    #         'name': 'Scale Secrets Store',
    #         'description': 'Secrets store for all secrets used by Scale',
    #         'driver': 'vault',
    #         'initialized': True,
    #         'sealed': False
    #     })
    #     create_mount = self._make_request('POST', url, data=data, catch_error=False)

    #     if create_mount.status_code == 403:
    #         raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
    #     elif create_mount.status_code not in [201, 409]:
    #         raise InvalidSecretsRequest('Invalid request return: ' + self.secrets_error_codes[r.status_code])

    def _create_vault_store(self):
        """Create a new store within the Vault backend that will be used for all Scale secrets.
        """

        url = self.secrets_url + '/sys/mounts/scale'
        headers = {
            "Content-Type": "application/json",
            "X-Vault-Token": self.secrets_token
        }
        data = json.dumps({
            'type': 'generic',
            'description': 'Secrets store for all secrets used by Scale'
        })

        create_mount = self._make_request('POST', url, headers=headers, data=data)

    def _dcos_authenticate(self):
        """Authenticate with DC/OS Vault backend and expect a status code 200 returned.
        """

        try:
            token = jwt.encode({'uid': self.service_account}, self.secrets_token, algorithm='RS256')
        except ValueError:
            raise InvalidSecretsToken('The provided token could not be encoded: ')

        url = self.secrets_url + '/acs/api/v1/auth/login'
        data = json.dumps({
            'uid': self.service_account, 'token':token
        })

        request_auth = self._make_request('GET', url, data=data)

        self.secrets_url += '/secrets/v1/'
        access_token = [k + '=' + v for k, v in request_auth.json().items()]

        return access_token

    def _make_request(self, method, url, headers=None, data=None, catch_error=True):
        """Make a request to the secrets backend with the provided variables

        :param method: string that determines GET or POST request type
        :type method: str
        :param url: string containing the url for the request
        :type url: str
        :param headers: headers to attach to the request
        :type headers: json
        :param data: data to attach to the request
        :type data: json
        :param data: flag to catch error before return
        :type data: bool

        :return: an object containing information from the request
        :rtype: requests.request
        """

        if not headers:
            headers = {
                'Content-Type': 'application/json'
            }

        if not data:
            data = {}

        r = requests.request(method=method, url=url, headers=headers, data=data)

        if catch_error and r.status_code in self.secrets_error_codes:
            if r.status_code == 403:
                raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
            else:
                raise InvalidSecretsRequest('Invalid request return: ' + self.secrets_error_codes[r.status_code])

        return r

    def _vault_authenticate(self):
        """Authenticate with Vault and expect a status code 200 returned.
        """

        url = self.secrets_url + '/v1/sys/health'
        request_auth = self._make_request('GET', url)
        self.secrets_url += '/v1'
<commit_msg>bog fixes: secrets handler<commit_after>"""Handles secret getters and setters for Scale"""

import ast
import jwt
import json
import requests

from django.conf import settings
from vault.exceptions import InvalidSecretsAuthorization, InvalidSecretsRequest, InvalidSecretsToken, InvalidSecretsValue

class SecretsHandler(object):
    """Represents a secrets handler for setting and retrieving secrets
    """

    def __init__(self):
        """Creates a secrets handler object.  The backend is initially tested to ensure it exists and Scale can
        authenticate properly with it.
        """

        self.secrets_error_codes = {
            301: 'Status Code 301 - Too many "/" separators in URL.',
            400: 'Status Code 400 - Invalid request, missing or invalid data.',
            403: 'Status Code 403 - Forbidden, authentication details are incorrect or no access to feature.',
            404: 'Status Code 404 - Invalid path, secret or backend not found.',
            405: 'Status Code 405 - Unsupported operation.',
            409: 'Status Code 409 - Secret or secret store already exists.',
            429: 'Status Code 429 - Rate limit exceeded.',
            500: 'Status Code 500 - Internal server error.',
            503: 'Status Code 503 - Secrets store is down for maintenance or is currently sealed.',
        }

        self.secrets_url = settings.SECRETS_URL
        self.secrets_token = settings.SECRETS_TOKEN
        self.service_account = settings.DCOS_SERVICE_ACCOUNT

        if self.service_account:
            self.dcos_token = self._dcos_authenticate()
        else:
            self.dcos_token = None
            self._vault_authenticate()

        self._check_secrets_backend()

    def get_job_type_secrets(self, job_name):
        """Retrieves the secrets located at the job_name within the backend

        :param job_name: path within the secrets backend that the secret is stored
        :type job_name: str
        :return: secret_values
        :rtype: str
        """

        url = self.secrets_url

        if self.dcos_token:
            url = ''.join([url, '/secret/default/scale/job-type/', job_name])
            data = json.dumps({'uid': self.service_account, 'token': self.dcos_token})
            get_secret = self._make_request('GET', url, data=data)

            try:
                response = get_secret.json()
                secret_values = ast.literal_eval(response['value'])
            except SyntaxError:
                raise InvalidSecretsValue('The DCOS secrets value could not be converted to JSON.')

        else:
            url = url + '/secret/scale/job-type/' + job_name
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            get_secret = self._make_request('GET', url, headers=headers)

            response = get_secret.json()
            secret_values = response['data']

        return secret_values

    def list_job_types(self):
        """Gets the names of all job types that have secrets

        :return: all_job_types
        :rtype: list
        """

        url = self.secrets_url

        if self.dcos_token:
            url += '/secret/default/scale/job-type?list=true'
            data = json.dumps({'uid': self.service_account, 'token': self.dcos_token})
            list_jobs = self._make_request('GET', url, data=data)

            response = list_jobs.json()
            all_job_types = response['array']

        else:
            url += '/secret/scale/job-type'
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            list_jobs = self._make_request('LIST', url, headers)

            response = list_jobs.json()
            all_job_types = response['data']['keys']

        return all_job_types

    def set_job_type_secrets(self, job_name, secret_json):
        """write job-type secrets to the secrets backend

        :param job_name: name of the job that the screts belong to.
        :type job_name: str
        :param secret_json: dict with name:value pairs for all secrets associated with the job
        :type secret_json: dict
        :return:
        """

        url = self.secrets_url
        secret_values = json.dumps(secret_json)

        if self.dcos_token:
            url = ''.join([url, '/secret/default/scale/job-type/', job_name])
            data = json.dumps({
                'uid': self.service_account,
                'token': self.dcos_token,
                'author': 'scale',
                'value': secret_values
            })
            set_secret = self._make_request('PUT', url, data=data)

        else:
            url = ''.join([url, '/secret/scale/job-type/', job_name])
            headers = {
                'Content-Type': 'application/json',
                'X-Vault-Token': self.secrets_token
            }
            data = secret_values
            set_secret = self._make_request('PUT', url, headers, data)

    def _check_secrets_backend(self):
        """Validates that Scale can transact with the secrets backend properly.
        """

        url = self.secrets_url

        if self.dcos_token:
            pass

        else:
            url += '/sys/mounts'
            headers = {
                "Content-Type": "application/json",
                "X-Vault-Token": self.secrets_token
            }

            check_mount = self._make_request('GET', url, headers=headers, catch_error=False)

            if check_mount.status_code == 200:
                if 'scale/' not in json.loads(check_mount.content).keys():
                    self._create_vault_store()
            elif check_mount.status_code == 403:
                raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
            elif check_mount.status_code == 404:
                self._create_dcos_store()

    # def _create_dcos_store(self):
    #     """Create a new store within the DC/OS Vault backend that will be used for all Scale secrets.

    #     **As of 9FEB2017 DC/OS does not support additional backends.**
    #     """

    #     url = self.secrets_url + '/store/scale'
    #     data = json.dumps({
    #         'uid': self.service_account,
    #         'token': self.dcos_token,
    #         'name': 'Scale Secrets Store',
    #         'description': 'Secrets store for all secrets used by Scale',
    #         'driver': 'vault',
    #         'initialized': True,
    #         'sealed': False
    #     })
    #     create_mount = self._make_request('POST', url, data=data, catch_error=False)

    #     if create_mount.status_code == 403:
    #         raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
    #     elif create_mount.status_code not in [201, 409]:
    #         raise InvalidSecretsRequest('Invalid request return: ' + self.secrets_error_codes[r.status_code])

    def _create_vault_store(self):
        """Create a new store within the Vault backend that will be used for all Scale secrets.
        """

        url = self.secrets_url + '/sys/mounts/scale'
        headers = {
            "Content-Type": "application/json",
            "X-Vault-Token": self.secrets_token
        }
        data = json.dumps({
            'type': 'generic',
            'description': 'Secrets store for all secrets used by Scale'
        })

        create_mount = self._make_request('POST', url, headers=headers, data=data)

        # A store needs at least one secret in it for other functions to work...
        backends = ['/job-type/placeholder', '/internal/placeholder', '/storage/placeholder']
        data = json.dumps({'empty': 'secret'})

        for folder in backends:
            url = self.secrets_url + '/secret/scale' + folder
            set_secret = self._make_request('PUT', url, headers, data)

    def _dcos_authenticate(self):
        """Authenticate with DC/OS Vault backend and expect a status code 200 returned.
        """

        try:
            token = jwt.encode({'uid': self.service_account}, self.secrets_token, algorithm='RS256')
        except ValueError:
            raise InvalidSecretsToken('The provided token could not be encoded: ')

        url = self.secrets_url + '/acs/api/v1/auth/login'
        data = json.dumps({
            'uid': self.service_account, 'token':token
        })

        request_auth = self._make_request('GET', url, data=data)

        self.secrets_url += '/secrets/v1'
        access_token = [k + '=' + v for k, v in request_auth.json().items()]

        return access_token

    def _make_request(self, method, url, headers=None, data=None, catch_error=True):
        """Make a request to the secrets backend with the provided variables

        :param method: string that determines GET or POST request type
        :type method: str
        :param url: string containing the url for the request
        :type url: str
        :param headers: headers to attach to the request
        :type headers: json
        :param data: data to attach to the request
        :type data: json
        :param data: flag to catch error before return
        :type data: bool

        :return: an object containing information from the request
        :rtype: requests.request
        """

        if not headers:
            headers = {
                'Content-Type': 'application/json'
            }

        if not data:
            data = {}

        r = requests.request(method=method, url=url, headers=headers, data=data)

        if catch_error and r.status_code in self.secrets_error_codes:
            if r.status_code == 403:
                raise InvalidSecretsAuthorization('Permission was denied when getting a secret')
            else:
                raise InvalidSecretsRequest('Invalid request return: ' + self.secrets_error_codes[r.status_code])

        return r

    def _vault_authenticate(self):
        """Authenticate with Vault and expect a status code 200 returned.
        """

        url = self.secrets_url + '/v1/sys/health'
        request_auth = self._make_request('GET', url)
        self.secrets_url += '/v1'
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
This code demonstrates how to use dedupe with a comma separated values
(CSV) file. All operations are performed in memory, so will run very
quickly on datasets up to ~10,000 rows.

We start with a CSV file containing our messy data. In this example,
it is listings of early childhood education centers in Chicago
compiled from several different sources.

The output will be a CSV with our clustered results.

For larger datasets, see our [mysql_example](http://open-city.github.com/dedupe/doc/mysql_example.html)
"""

import os
import csv
import re
import collections
import logging
import optparse
import numpy as np
from numpy import nan
from numpy import matlib

import dedupe
import dedupe.centroid

# ## Logging

# Dedupe uses Python logging to show or suppress verbose output. Added for convenience.
# To enable verbose logging, run `python examples/csv_example/csv_example.py -v`
optp = optparse.OptionParser()
optp.add_option('-v', '--verbose', dest='verbose', action='count',
                help='Increase verbosity (specify multiple times for more)'
                )
(opts, args) = optp.parse_args()
log_level = logging.WARNING
if opts.verbose == 1:
    log_level = logging.INFO
elif opts.verbose >= 2:
    log_level = logging.DEBUG
logging.getLogger().setLevel(log_level)

# ## Setup

# Switch to our working directory and set up our input and out put paths,
# as well as our settings and training file locations
input_file = 'csv_example_messy_input.csv'
output_file = 'csv_example_output.csv'
settings_file = 'csv_example_learned_settings'
training_file = 'csv_example_training.json'

# Dedupe can take custom field comparison functions, here's one
# we'll use for zipcodes
def sameOrNotComparator(field_1, field_2) :
    if field_1 and field_2 :
        if field_1 == field_2 :
            return 1
        else:
            return 0
    else :
        return nan

def preProcess(column):
    """
    Do a little bit of data cleaning with the help of
    [AsciiDammit](https://github.com/tnajdek/ASCII--Dammit) and
    Regex. Things like casing, extra spaces, quotes and new lines can
    be ignored.
    """

    column = dedupe.asciiDammit(column)
    column = re.sub('  +', ' ', column)
    column = re.sub('\n', ' ', column)
    column = column.strip().strip('"').strip("'").lower().strip()
    return column

def readData(filename):
    """
    Read in our data from a CSV file and create a dictionary of records,
    where the key is a unique record ID and each value is dict
    """

    data_d = {}
    with open(filename) as f:
        reader = csv.DictReader(f)
        for row in reader:
            clean_row = [(k, preProcess(v)) for (k, v) in row.items()]
            row_id = int(row['Id'])
            data_d[row_id] = dict(clean_row)

    return data_d

print 'importing data ...'
data_d = readData(input_file)

# ## Training

if os.path.exists(settings_file):
    print 'reading from', settings_file
    with open(settings_file, 'rb') as f:
        deduper = dedupe.StaticDedupe(f)

else:
    # Define the fields dedupe will pay attention to
    #
    # Notice how we are telling dedupe to use a custom field comparator
    # for the 'Zip' field.
    fields = {
        'Site name': {'type': 'String'},
        'Address': {'type': 'String'},
        'Zip': {'type': 'Custom',
                'comparator' : sameOrNotComparator,
                'Has Missing' : True},
        'Phone': {'type': 'String', 'Has Missing' : True},
        }

    # Create a new deduper object and pass our data model to it.
    deduper = dedupe.Dedupe(fields)

    # To train dedupe, we feed it a random sample of records.
    deduper.sample(data_d, 150000)

    # If we have training data saved from a previous run of dedupe,
    # look for it an load it in.
    # __Note:__ if you want to train from scratch, delete the training_file
    if os.path.exists(training_file):
        print 'reading labeled examples from ', training_file
        with open(training_file, 'rb') as f:
            deduper.readTraining(f)

    # ## Active learning
    # Dedupe will find the next pair of records
    # it is least certain about and ask you to label them as duplicates
    # or not.
    # use 'y', 'n' and 'u' keys to flag duplicates
    # press 'f' when you are finished
    print 'starting active labeling...'

    dedupe.consoleLabel(deduper)

    deduper.train()

    # When finished, save our training away to disk
    deduper.writeTraining(training_file)

    # Save our weights and predicates to disk.  If the settings file
    # exists, we will skip all the training and learning next time we run
    # this file.
    deduper.writeSettings(settings_file)

# ## Blocking

print 'blocking...'

# ## Clustering

# Find the threshold that will maximize a weighted average of our precision and recall.
# When we set the recall weight to 2, we are saying we care twice as much
# about recall as we do precision.
#
# If we had more data, we would not pass in all the blocked data into
# this function but a representative sample.

threshold = deduper.threshold(data_d, recall_weight=2)

# `match` will return sets of record IDs that dedupe
# believes are all referring to the same entity.

print 'clustering...'
clustered_dupes = deduper.match(data_d, threshold)

print '# duplicate sets', len(clustered_dupes)

# ## Writing Results

# Write our original data back out to a CSV with a new column called
# 'Cluster ID' which indicates which records refer to each other.

cluster_membership = {}
for (cluster_id, cluster) in enumerate(clustered_dupes):
    id_set, conf_score = cluster
    cluster_d = [data_d[c] for c in id_set]
    canonical_rep = dedupe.canonicalize(cluster_d)
    print canonical_rep
    for record_id in id_set:
        cluster_membership[record_id] = {
            "cluster id" : cluster_id,
            "canonical representation" : canonical_rep,
            "confidence": conf_score
        }

with open(output_file, 'w') as f_output:
    writer = csv.writer(f_output)

    with open(input_file) as f_input :
        reader = csv.reader(f_input)

        heading_row = reader.next()
        heading_row.insert(0, 'Cluster ID')
        canonical_keys = canonical_rep.keys()
        for key in canonical_keys:
            heading_row.append('canonical_' + key)
        heading_row.append('confidence_score')

        writer.writerow(heading_row)

        for row in reader:
            row_id = int(row[0])
            if row_id in cluster_membership:
                cluster_id = cluster_membership[row_id]["cluster id"]
                canonical_rep = cluster_membership[row_id]["canonical representation"]
                row.insert(0, cluster_id)
                for key in canonical_keys:
                    row.append(canonical_rep[key])
                row.append(cluster_membership[row_id]['confidence'])
            else:
                row.insert(0, None)
                for key in canonical_keys:
                    row.append(None)
                row.append(None)
            writer.writerow(row)
<commit_msg>removed unnecessary imports<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
This code demonstrates how to use dedupe with a comma separated values
(CSV) file. All operations are performed in memory, so will run very
quickly on datasets up to ~10,000 rows.

We start with a CSV file containing our messy data. In this example,
it is listings of early childhood education centers in Chicago
compiled from several different sources.

The output will be a CSV with our clustered results.

For larger datasets, see our [mysql_example](http://open-city.github.com/dedupe/doc/mysql_example.html)
"""

import os
import csv
import re
import collections
import logging
import optparse
import numpy as np
from numpy import nan

import dedupe

# ## Logging

# Dedupe uses Python logging to show or suppress verbose output. Added for convenience.
# To enable verbose logging, run `python examples/csv_example/csv_example.py -v`
optp = optparse.OptionParser()
optp.add_option('-v', '--verbose', dest='verbose', action='count',
                help='Increase verbosity (specify multiple times for more)'
                )
(opts, args) = optp.parse_args()
log_level = logging.WARNING
if opts.verbose == 1:
    log_level = logging.INFO
elif opts.verbose >= 2:
    log_level = logging.DEBUG
logging.getLogger().setLevel(log_level)

# ## Setup

# Switch to our working directory and set up our input and out put paths,
# as well as our settings and training file locations
input_file = 'csv_example_messy_input.csv'
output_file = 'csv_example_output.csv'
settings_file = 'csv_example_learned_settings'
training_file = 'csv_example_training.json'

# Dedupe can take custom field comparison functions, here's one
# we'll use for zipcodes
def sameOrNotComparator(field_1, field_2) :
    if field_1 and field_2 :
        if field_1 == field_2 :
            return 1
        else:
            return 0
    else :
        return nan

def preProcess(column):
    """
    Do a little bit of data cleaning with the help of
    [AsciiDammit](https://github.com/tnajdek/ASCII--Dammit) and
    Regex. Things like casing, extra spaces, quotes and new lines can
    be ignored.
    """

    column = dedupe.asciiDammit(column)
    column = re.sub('  +', ' ', column)
    column = re.sub('\n', ' ', column)
    column = column.strip().strip('"').strip("'").lower().strip()
    return column

def readData(filename):
    """
    Read in our data from a CSV file and create a dictionary of records,
    where the key is a unique record ID and each value is dict
    """

    data_d = {}
    with open(filename) as f:
        reader = csv.DictReader(f)
        for row in reader:
            clean_row = [(k, preProcess(v)) for (k, v) in row.items()]
            row_id = int(row['Id'])
            data_d[row_id] = dict(clean_row)

    return data_d

print 'importing data ...'
data_d = readData(input_file)

# ## Training

if os.path.exists(settings_file):
    print 'reading from', settings_file
    with open(settings_file, 'rb') as f:
        deduper = dedupe.StaticDedupe(f)

else:
    # Define the fields dedupe will pay attention to
    #
    # Notice how we are telling dedupe to use a custom field comparator
    # for the 'Zip' field.
    fields = {
        'Site name': {'type': 'String'},
        'Address': {'type': 'String'},
        'Zip': {'type': 'Custom',
                'comparator' : sameOrNotComparator,
                'Has Missing' : True},
        'Phone': {'type': 'String', 'Has Missing' : True},
        }

    # Create a new deduper object and pass our data model to it.
    deduper = dedupe.Dedupe(fields)

    # To train dedupe, we feed it a random sample of records.
    deduper.sample(data_d, 150000)

    # If we have training data saved from a previous run of dedupe,
    # look for it an load it in.
    # __Note:__ if you want to train from scratch, delete the training_file
    if os.path.exists(training_file):
        print 'reading labeled examples from ', training_file
        with open(training_file, 'rb') as f:
            deduper.readTraining(f)

    # ## Active learning
    # Dedupe will find the next pair of records
    # it is least certain about and ask you to label them as duplicates
    # or not.
    # use 'y', 'n' and 'u' keys to flag duplicates
    # press 'f' when you are finished
    print 'starting active labeling...'

    dedupe.consoleLabel(deduper)

    deduper.train()

    # When finished, save our training away to disk
    deduper.writeTraining(training_file)

    # Save our weights and predicates to disk.  If the settings file
    # exists, we will skip all the training and learning next time we run
    # this file.
    deduper.writeSettings(settings_file)

# ## Blocking

print 'blocking...'

# ## Clustering

# Find the threshold that will maximize a weighted average of our precision and recall.
# When we set the recall weight to 2, we are saying we care twice as much
# about recall as we do precision.
#
# If we had more data, we would not pass in all the blocked data into
# this function but a representative sample.

threshold = deduper.threshold(data_d, recall_weight=2)

# `match` will return sets of record IDs that dedupe
# believes are all referring to the same entity.

print 'clustering...'
clustered_dupes = deduper.match(data_d, threshold)

print '# duplicate sets', len(clustered_dupes)

# ## Writing Results

# Write our original data back out to a CSV with a new column called
# 'Cluster ID' which indicates which records refer to each other.

cluster_membership = {}
for (cluster_id, cluster) in enumerate(clustered_dupes):
    id_set, conf_score = cluster
    cluster_d = [data_d[c] for c in id_set]
    canonical_rep = dedupe.canonicalize(cluster_d)
    print canonical_rep
    for record_id in id_set:
        cluster_membership[record_id] = {
            "cluster id" : cluster_id,
            "canonical representation" : canonical_rep,
            "confidence": conf_score
        }

with open(output_file, 'w') as f_output:
    writer = csv.writer(f_output)

    with open(input_file) as f_input :
        reader = csv.reader(f_input)

        heading_row = reader.next()
        heading_row.insert(0, 'Cluster ID')
        canonical_keys = canonical_rep.keys()
        for key in canonical_keys:
            heading_row.append('canonical_' + key)
        heading_row.append('confidence_score')

        writer.writerow(heading_row)

        for row in reader:
            row_id = int(row[0])
            if row_id in cluster_membership:
                cluster_id = cluster_membership[row_id]["cluster id"]
                canonical_rep = cluster_membership[row_id]["canonical representation"]
                row.insert(0, cluster_id)
                for key in canonical_keys:
                    row.append(canonical_rep[key])
                row.append(cluster_membership[row_id]['confidence'])
            else:
                row.insert(0, None)
                for key in canonical_keys:
                    row.append(None)
                row.append(None)
            writer.writerow(row)
<|endoftext|>