<commit_before>from pypy.rlib.objectmodel import newlist_hint, compute_hash
from pypy.rlib.rarithmetic import intmask
from pypy.rlib.rerased import new_static_erasing_pair

from rupypy.module import ClassDef
from rupypy.modules.comparable import Comparable
from rupypy.objects.objectobject import W_Object

def create_trans_table(source, replacement, inv=False):
    src = expand_trans_str(source, len(source), inv)
    repl = expand_trans_str(replacement, len(src))
    table = [chr(i) for i in xrange(256)]
    for i, c in enumerate(src):
        table[ord(c)] = repl[i]
    return table

def expand_trans_str(source, res_len, inv=False):
    # check the source for range definitions
    # and insert the missing characters
    expanded_source = []
    char = ""
    for i in range(res_len):
        if i < len(source):
            char = source[i]
        if char == "-":
            # expand the range
            assert 0 < i < len(source) - 1
            range_beg = ord(source[i - 1])
            range_end = ord(source[i + 1])
            for j in range(range_beg + 1, range_end - 1):
                expanded_source.append(chr(j))
        elif char:
            expanded_source.append(char[0])

    if inv:
        inverted_source = []
        # invert the source
        for i in range(256):
            if chr(i) not in expanded_source:
                inverted_source.append(chr(i))
        return inverted_source

    return expanded_source

class StringStrategy(object):
    def __init__(self, space):
        pass

class ConstantStringStrategy(StringStrategy):
    erase, unerase = new_static_erasing_pair("constant")

    def str_w(self, storage):
        return self.unerase(storage)

    def liststr_w(self, storage):
        strvalue = self.unerase(storage)
        return [c for c in strvalue]

    def length(self, storage):
        return len(self.unerase(storage))

    def getitem(self, storage, idx):
        return self.unerase(storage)[idx]

    def getslice(self, space, storage, start, end):
        return space.newstr_fromstr(self.unerase(storage)[start:end])

    def hash(self, storage):
        return compute_hash(self.unerase(storage))

    def copy(self, space, storage):
        return W_StringObject(space, storage, self)

    def to_mutable(self, space, s):
        s.strategy = strategy = space.fromcache(MutableStringStrategy)
        s.str_storage = strategy.erase(self.liststr_w(s.str_storage))

    def extend_into(self, src_storage, dst_storage):
        dst_storage += self.unerase(src_storage)

class MutableStringStrategy(StringStrategy):
    erase, unerase = new_static_erasing_pair("mutable")

    def str_w(self, storage):
        return "".join(self.unerase(storage))

    def liststr_w(self, storage):
        return self.unerase(storage)

    def length(self, storage):
        return len(self.unerase(storage))

    def getitem(self, storage, idx):
        return self.unerase(storage)[idx]

    def getslice(self, space, storage, start, end):
        return space.newstr_fromchars(self.unerase(storage)[start:end])

    def hash(self, storage):
        storage = self.unerase(storage)
        length = len(storage)
        if length == 0:
            return -1
        x = ord(storage[0]) << 7
        i = 0
        while i < length:
            x = intmask((1000003 * x) ^ ord(storage[i]))
            i += 1
        x ^= length
        return intmask(x)

    def copy(self, space, storage):
        return W_StringObject(space, storage, self)

    def to_mutable(self, space, s):
        pass

    def extend_into(self, src_storage, dst_storage):
        dst_storage += self.unerase(src_storage)

    def clear(self, s):
        storage = self.unerase(s.str_storage)
        del storage[:]

    def downcase(self, storage):
        storage = self.unerase(storage)
        changed = False
        for i, c in enumerate(storage):
            # TODO: obscure hack because lower() returns a string, rather than
            # a char, this should be fixed upstream.
            new_c = c.lower()[0]
            changed |= (c != new_c)
            storage[i] = new_c
        return changed

class W_StringObject(W_Object):
    classdef = ClassDef("String", W_Object.classdef)
    classdef.include_module(Comparable)

    def __init__(self, space, storage, strategy):
        W_Object.__init__(self, space)
        self.str_storage = storage
        self.strategy = strategy

    @staticmethod
    def newstr_fromstr(space, strvalue):
        strategy = space.fromcache(ConstantStringStrategy)
        storage = strategy.erase(strvalue)
        return W_StringObject(space, storage, strategy)

    @staticmethod
    def newstr_fromchars(space, chars):
        strategy = space.fromcache(MutableStringStrategy)
        storage = strategy.erase(chars)
        return W_StringObject(space, storage, strategy)

    def str_w(self, space):
        return self.strategy.str_w(self.str_storage)

    def liststr_w(self, space):
        return self.strategy.liststr_w(self.str_storage)

    def length(self):
        return self.strategy.length(self.str_storage)

    def copy(self, space):
        return self.strategy.copy(space, self.str_storage)

    def replace(self, space, chars):
        strategy = space.fromcache(MutableStringStrategy)
        self.str_storage = strategy.erase(chars)
        self.strategy = strategy

    def extend(self, space, w_other):
        self.strategy.to_mutable(space, self)
        strategy = self.strategy
        assert isinstance(strategy, MutableStringStrategy)
        storage = strategy.unerase(self.str_storage)
        w_other.strategy.extend_into(w_other.str_storage, storage)

    def clear(self, space):
        self.strategy.to_mutable(space, self)
        self.strategy.clear(self)

    def tr_trans(self, space, source, replacement, squeeze):
        change_made = False
        string = space.str_w(self)
        new_string = []
        is_negative_set = len(source) > 1 and source[0] == "^"
        if is_negative_set:
            source = source[1:]

        trans_table = create_trans_table(source, replacement, is_negative_set)

        if squeeze:
            last_repl = ""
            for char in string:
                repl = trans_table[ord(char)]
                if last_repl == repl:
                    continue
                if repl != char:
                    last_repl = repl
                    if not change_made:
                        change_made = True
                new_string.append(repl)
        else:
            for char in string:
                repl = trans_table[ord(char)]
                if not change_made and repl != char:
                    change_made = True
                new_string.append(repl)

        return new_string if change_made else None

    @classdef.method("to_str")
    @classdef.method("to_s")
    def method_to_s(self, space):
        return self

    @classdef.method("+")
    def method_plus(self, space, w_other):
        assert isinstance(w_other, W_StringObject)
        total_size = self.length() + w_other.length()
        s = space.newstr_fromchars(newlist_hint(total_size))
        s.extend(space, self)
        s.extend(space, w_other)
        return s

    @classdef.method("<<")
    def method_lshift(self, space, w_other):
        assert isinstance(w_other, W_StringObject)
        self.extend(space, w_other)
        return self

    @classdef.method("size")
    @classdef.method("length")
    def method_length(self, space):
        return space.newint(self.length())

    @classdef.method("hash")
    def method_hash(self, space):
        return space.newint(self.strategy.hash(self.str_storage))

    @classdef.method("[]")
    def method_subscript(self, space, w_idx, w_count=None):
        start, end, as_range, nil = space.subscript_access(self.length(), w_idx, w_count=w_count)
        if nil:
            return space.w_nil
        elif as_range:
            assert start > 0
            return self.strategy.getslice(space, self.str_storage, start, end)
        else:
            return space.newstr_fromstr(self.strategy.getitem(self.str_storage, start))

    @classdef.method("<=>")
    def method_comparator(self, space, w_other):
        if isinstance(w_other, W_StringObject):
            s1 = space.str_w(self)
            s2 = space.str_w(w_other)
            if s1 < s2:
                return space.newint(-1)
            elif s1 == s2:
                return space.newint(0)
            elif s1 > s2:
                return space.newint(1)
        else:
            if space.respond_to(w_other, space.newsymbol("to_str")) and space.respond_to(w_other, space.newsymbol("<=>")):
                tmp = space.send(w_other, space.newsymbol("<=>"), [self])
                if tmp is not space.w_nil:
                    return space.newint(-space.int_w(tmp))
            return space.w_nil

    @classdef.method("freeze")
    def method_freeze(self, space):
        pass

    @classdef.method("dup")
    def method_dup(self, space):
        return self.copy(space)

    @classdef.method("to_sym")
    @classdef.method("intern")
    def method_to_sym(self, space):
        return space.newsymbol(space.str_w(self))

    @classdef.method("clear")
    def method_clear(self, space):
        self.clear(space)
        return self

    @classdef.method("ljust", integer="int", padstr="str")
    def method_ljust(self, space, integer, padstr=" "):
        if not padstr:
            raise space.error(space.w_ArgumentError, "zero width padding")
        elif integer <= self.length():
            return self.copy(space)
        else:
            pad_len = integer - self.length() - 1
            assert pad_len >= 0
            chars = []
            chars += space.str_w(self)
            for i in xrange(pad_len / len(padstr)):
                chars += padstr
            chars += padstr[:pad_len % len(padstr) + 1]
            return space.newstr_fromchars(chars)

    @classdef.method("split", limit="int")
    def method_split(self, space, w_sep=None, limit=-1):
        if w_sep is None:
            sep = None
        elif isinstance(w_sep, W_StringObject):
            sep = space.str_w(w_sep)
        else:
            raise NotImplementedError("Regexp separators for String#split")
        results = space.str_w(self).split(sep, limit - 1)
        return space.newarray([space.newstr_fromstr(s) for s in results])

    @classdef.method("downcase!")
    def method_downcase(self, space):
        self.strategy.to_mutable(space, self)
        changed = self.strategy.downcase(self.str_storage)
        return self if changed else space.w_nil

    @classdef.method("to_i", radix="int")
    def method_to_i(self, space, radix=10):
        return space.newint(int(space.str_w(self), radix))

    @classdef.method("tr", source="str", replacement="str")
    def method_tr(self, space, source, replacement):
        string = self.copy(space)
        new_string = self.tr_trans(space, source, replacement, False)
        return space.newstr_fromchars(new_string) if new_string else string

    @classdef.method("tr!", source="str", replacement="str")
    def method_tr_i(self, space, source, replacement):
        new_string = self.tr_trans(space, source, replacement, False)
        self.replace(space, new_string)
        return self if new_string else space.w_nil

    @classdef.method("tr_s", source="str", replacement="str")
    def method_tr_s(self, space, source, replacement):
        string = self.copy(space)
        new_string = self.tr_trans(space, source, replacement, True)
        return space.newstr_fromchars(new_string) if new_string else string

    @classdef.method("tr_s!", source="str", replacement="str")
    def method_tr_s_i(self, space, source, replacement):
        new_string = self.tr_trans(space, source, replacement, True)
        self.replace(space, new_string)
        return self if new_string else space.w_nil
<commit_msg>more accurate assert<commit_after>from pypy.rlib.objectmodel import newlist_hint, compute_hash
from pypy.rlib.rarithmetic import intmask
from pypy.rlib.rerased import new_static_erasing_pair

from rupypy.module import ClassDef
from rupypy.modules.comparable import Comparable
from rupypy.objects.objectobject import W_Object

def create_trans_table(source, replacement, inv=False):
    src = expand_trans_str(source, len(source), inv)
    repl = expand_trans_str(replacement, len(src))
    table = [chr(i) for i in xrange(256)]
    for i, c in enumerate(src):
        table[ord(c)] = repl[i]
    return table

def expand_trans_str(source, res_len, inv=False):
    # check the source for range definitions
    # and insert the missing characters
    expanded_source = []
    char = ""
    for i in range(res_len):
        if i < len(source):
            char = source[i]
        if char == "-":
            # expand the range
            assert 0 < i < len(source) - 1
            range_beg = ord(source[i - 1])
            range_end = ord(source[i + 1])
            for j in range(range_beg + 1, range_end - 1):
                expanded_source.append(chr(j))
        elif char:
            expanded_source.append(char[0])

    if inv:
        inverted_source = []
        # invert the source
        for i in range(256):
            if chr(i) not in expanded_source:
                inverted_source.append(chr(i))
        return inverted_source

    return expanded_source

class StringStrategy(object):
    def __init__(self, space):
        pass

class ConstantStringStrategy(StringStrategy):
    erase, unerase = new_static_erasing_pair("constant")

    def str_w(self, storage):
        return self.unerase(storage)

    def liststr_w(self, storage):
        strvalue = self.unerase(storage)
        return [c for c in strvalue]

    def length(self, storage):
        return len(self.unerase(storage))

    def getitem(self, storage, idx):
        return self.unerase(storage)[idx]

    def getslice(self, space, storage, start, end):
        return space.newstr_fromstr(self.unerase(storage)[start:end])

    def hash(self, storage):
        return compute_hash(self.unerase(storage))

    def copy(self, space, storage):
        return W_StringObject(space, storage, self)

    def to_mutable(self, space, s):
        s.strategy = strategy = space.fromcache(MutableStringStrategy)
        s.str_storage = strategy.erase(self.liststr_w(s.str_storage))

    def extend_into(self, src_storage, dst_storage):
        dst_storage += self.unerase(src_storage)

class MutableStringStrategy(StringStrategy):
    erase, unerase = new_static_erasing_pair("mutable")

    def str_w(self, storage):
        return "".join(self.unerase(storage))

    def liststr_w(self, storage):
        return self.unerase(storage)

    def length(self, storage):
        return len(self.unerase(storage))

    def getitem(self, storage, idx):
        return self.unerase(storage)[idx]

    def getslice(self, space, storage, start, end):
        return space.newstr_fromchars(self.unerase(storage)[start:end])

    def hash(self, storage):
        storage = self.unerase(storage)
        length = len(storage)
        if length == 0:
            return -1
        x = ord(storage[0]) << 7
        i = 0
        while i < length:
            x = intmask((1000003 * x) ^ ord(storage[i]))
            i += 1
        x ^= length
        return intmask(x)

    def copy(self, space, storage):
        return W_StringObject(space, storage, self)

    def to_mutable(self, space, s):
        pass

    def extend_into(self, src_storage, dst_storage):
        dst_storage += self.unerase(src_storage)

    def clear(self, s):
        storage = self.unerase(s.str_storage)
        del storage[:]

    def downcase(self, storage):
        storage = self.unerase(storage)
        changed = False
        for i, c in enumerate(storage):
            # TODO: obscure hack because lower() returns a string, rather than
            # a char, this should be fixed upstream.
            new_c = c.lower()[0]
            changed |= (c != new_c)
            storage[i] = new_c
        return changed

class W_StringObject(W_Object):
    classdef = ClassDef("String", W_Object.classdef)
    classdef.include_module(Comparable)

    def __init__(self, space, storage, strategy):
        W_Object.__init__(self, space)
        self.str_storage = storage
        self.strategy = strategy

    @staticmethod
    def newstr_fromstr(space, strvalue):
        strategy = space.fromcache(ConstantStringStrategy)
        storage = strategy.erase(strvalue)
        return W_StringObject(space, storage, strategy)

    @staticmethod
    def newstr_fromchars(space, chars):
        strategy = space.fromcache(MutableStringStrategy)
        storage = strategy.erase(chars)
        return W_StringObject(space, storage, strategy)

    def str_w(self, space):
        return self.strategy.str_w(self.str_storage)

    def liststr_w(self, space):
        return self.strategy.liststr_w(self.str_storage)

    def length(self):
        return self.strategy.length(self.str_storage)

    def copy(self, space):
        return self.strategy.copy(space, self.str_storage)

    def replace(self, space, chars):
        strategy = space.fromcache(MutableStringStrategy)
        self.str_storage = strategy.erase(chars)
        self.strategy = strategy

    def extend(self, space, w_other):
        self.strategy.to_mutable(space, self)
        strategy = self.strategy
        assert isinstance(strategy, MutableStringStrategy)
        storage = strategy.unerase(self.str_storage)
        w_other.strategy.extend_into(w_other.str_storage, storage)

    def clear(self, space):
        self.strategy.to_mutable(space, self)
        self.strategy.clear(self)

    def tr_trans(self, space, source, replacement, squeeze):
        change_made = False
        string = space.str_w(self)
        new_string = []
        is_negative_set = len(source) > 1 and source[0] == "^"
        if is_negative_set:
            source = source[1:]

        trans_table = create_trans_table(source, replacement, is_negative_set)

        if squeeze:
            last_repl = ""
            for char in string:
                repl = trans_table[ord(char)]
                if last_repl == repl:
                    continue
                if repl != char:
                    last_repl = repl
                    if not change_made:
                        change_made = True
                new_string.append(repl)
        else:
            for char in string:
                repl = trans_table[ord(char)]
                if not change_made and repl != char:
                    change_made = True
                new_string.append(repl)

        return new_string if change_made else None

    @classdef.method("to_str")
    @classdef.method("to_s")
    def method_to_s(self, space):
        return self

    @classdef.method("+")
    def method_plus(self, space, w_other):
        assert isinstance(w_other, W_StringObject)
        total_size = self.length() + w_other.length()
        s = space.newstr_fromchars(newlist_hint(total_size))
        s.extend(space, self)
        s.extend(space, w_other)
        return s

    @classdef.method("<<")
    def method_lshift(self, space, w_other):
        assert isinstance(w_other, W_StringObject)
        self.extend(space, w_other)
        return self

    @classdef.method("size")
    @classdef.method("length")
    def method_length(self, space):
        return space.newint(self.length())

    @classdef.method("hash")
    def method_hash(self, space):
        return space.newint(self.strategy.hash(self.str_storage))

    @classdef.method("[]")
    def method_subscript(self, space, w_idx, w_count=None):
        start, end, as_range, nil = space.subscript_access(self.length(), w_idx, w_count=w_count)
        if nil:
            return space.w_nil
        elif as_range:
            assert start >= 0
            return self.strategy.getslice(space, self.str_storage, start, end)
        else:
            return space.newstr_fromstr(self.strategy.getitem(self.str_storage, start))

    @classdef.method("<=>")
    def method_comparator(self, space, w_other):
        if isinstance(w_other, W_StringObject):
            s1 = space.str_w(self)
            s2 = space.str_w(w_other)
            if s1 < s2:
                return space.newint(-1)
            elif s1 == s2:
                return space.newint(0)
            elif s1 > s2:
                return space.newint(1)
        else:
            if space.respond_to(w_other, space.newsymbol("to_str")) and space.respond_to(w_other, space.newsymbol("<=>")):
                tmp = space.send(w_other, space.newsymbol("<=>"), [self])
                if tmp is not space.w_nil:
                    return space.newint(-space.int_w(tmp))
            return space.w_nil

    @classdef.method("freeze")
    def method_freeze(self, space):
        pass

    @classdef.method("dup")
    def method_dup(self, space):
        return self.copy(space)

    @classdef.method("to_sym")
    @classdef.method("intern")
    def method_to_sym(self, space):
        return space.newsymbol(space.str_w(self))

    @classdef.method("clear")
    def method_clear(self, space):
        self.clear(space)
        return self

    @classdef.method("ljust", integer="int", padstr="str")
    def method_ljust(self, space, integer, padstr=" "):
        if not padstr:
            raise space.error(space.w_ArgumentError, "zero width padding")
        elif integer <= self.length():
            return self.copy(space)
        else:
            pad_len = integer - self.length() - 1
            assert pad_len >= 0
            chars = []
            chars += space.str_w(self)
            for i in xrange(pad_len / len(padstr)):
                chars += padstr
            chars += padstr[:pad_len % len(padstr) + 1]
            return space.newstr_fromchars(chars)

    @classdef.method("split", limit="int")
    def method_split(self, space, w_sep=None, limit=-1):
        if w_sep is None:
            sep = None
        elif isinstance(w_sep, W_StringObject):
            sep = space.str_w(w_sep)
        else:
            raise NotImplementedError("Regexp separators for String#split")
        results = space.str_w(self).split(sep, limit - 1)
        return space.newarray([space.newstr_fromstr(s) for s in results])

    @classdef.method("downcase!")
    def method_downcase(self, space):
        self.strategy.to_mutable(space, self)
        changed = self.strategy.downcase(self.str_storage)
        return self if changed else space.w_nil

    @classdef.method("to_i", radix="int")
    def method_to_i(self, space, radix=10):
        return space.newint(int(space.str_w(self), radix))

    @classdef.method("tr", source="str", replacement="str")
    def method_tr(self, space, source, replacement):
        string = self.copy(space)
        new_string = self.tr_trans(space, source, replacement, False)
        return space.newstr_fromchars(new_string) if new_string else string

    @classdef.method("tr!", source="str", replacement="str")
    def method_tr_i(self, space, source, replacement):
        new_string = self.tr_trans(space, source, replacement, False)
        self.replace(space, new_string)
        return self if new_string else space.w_nil

    @classdef.method("tr_s", source="str", replacement="str")
    def method_tr_s(self, space, source, replacement):
        string = self.copy(space)
        new_string = self.tr_trans(space, source, replacement, True)
        return space.newstr_fromchars(new_string) if new_string else string

    @classdef.method("tr_s!", source="str", replacement="str")
    def method_tr_s_i(self, space, source, replacement):
        new_string = self.tr_trans(space, source, replacement, True)
        self.replace(space, new_string)
        return self if new_string else space.w_nil
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function
import errno
import json
import os
import re
import shutil
import string
import sys

import requests

CONFIG_FILE = 'scripts.conf'
EXTENSIONS = ['avi', 'm4v', 'mkv', 'mp4']
SUB_EXTENSIONS = ['idx', 'sub', 'srt']
PATTERN = re.compile('^(.*)(\d{4})\.(.*)', re.I)
REPLACEMENTS = {
    ':': '',
}
ENDPOINT = 'https://api.themoviedb.org/3/search/movie'

def get_config(cfgfile):
    if not cfgfile.startswith('/'):
        cfgfile = os.path.join(os.path.dirname(sys.argv[0]), '..', cfgfile)

    config = {}
    with open(cfgfile) as fh:
        for line in fh:
            k, v = re.split('\s+=\s+', line, 1)
            config[k] = v.strip()
    return config

CONFIG = get_config(CONFIG_FILE)

def get_matching_files(dirname):
    # Loop through the files in ``job_dir``, trying to find one which
    # matches PATTERN. Also build a list of files and their sizes.
    #
    # If no files match PATTERN but we do find one or more files with
    # a suitable extension we assume the largest file is the one we
    # want. We take this file and use the ``job_dir`` as the base for
    # the new file name.
    files = {
        'match': (),
        'by_size': [],
        'subs': [],
    }
    for root, _, fns in os.walk(dirname):
        for fn in fns:
            ext = fn.rsplit('.', 1)[-1].lower()
            if ext in SUB_EXTENSIONS:
                files['subs'].append(fn)

            if ext not in EXTENSIONS:
                continue

            abs = os.path.abspath(os.path.join(root, fn))
            files['by_size'].append((fn, os.path.getsize(abs)))
            match = PATTERN.match(fn)
            if match:
                files['match'] = (abs, match)
                break

    files['by_size'] = sorted(files['by_size'], key=lambda x: x[1])

    return files

def get_canonical_name(name, year):
    """
    Translates a dotted filename into the real name with spaces while
    preserving acronyms.
    """
    if not name.endswith('.'):
        name += '.'

    regex = '((?:[A-Z]\.)+)'
    acronyms = re.findall(regex, name)
    parts = re.split(regex, name)

    if not parts:
        return name

    canonical_name = []
    for part in parts:
        if part in acronyms:
            part = re.sub('\.$', '', part)
        else:
            part = part.replace('.', ' ')
        canonical_name.append(part.strip())
    canonical_name = ' '.join(canonical_name)

    # look up TMDb for the show name
    canonical_name = get_tmdb_name(canonical_name, year)
    if not canonical_name:
        return None

    # replace trailing dots
    canonical_name = re.sub('\.+$', '', canonical_name)

    # replace unsuitable characters
    for k, v in REPLACEMENTS.items():
        canonical_name = canonical_name.replace(k, v)

    return canonical_name

def get_tmdb_name(name, year):
    headers = {'Accept': 'application/json'}
    params = {
        'api_key': CONFIG['TMDB_API_KEY'],
        'query': name,
    }
    response = requests.get(ENDPOINT, params=params, headers=headers)
    results = json.loads(response.content.decode('utf-8'))['results']
    if not results:
        return None

    for r in results:
        if r['release_date'].split('-')[0] == year:
            return r['original_title']

def mkdirp(path, mode=0755):
    try:
        os.makedirs(path, mode=mode)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def main(job_dir, category):
    files = get_matching_files(job_dir)

    if files['match']:
        fn, match = files['match']
    else:
        # no files at all? just exit
        if not files['by_size']:
            print("No suitable files found")
            sys.exit()

        # found a potential match
        fn = files['by_size'][0][0]
        fn = os.path.join(job_dir, fn)
        match = PATTERN.match(os.path.basename(job_dir))
        if not match:
            print("No suitable files found")
            sys.exit()

    dotted_name, year = match.groups()[:2]
    canonical_name = get_canonical_name(dotted_name, year)

    # bail out if TMBd doesn't find anything
    if not canonical_name:
        print("Could not find files in TMDb for %s %s" % (dotted_name, year))
        sys.exit()

    dotted_name = canonical_name.replace(' ', '.')

    base_name = "%s.%s" % (dotted_name, year)
    for fn in [fn] + files['subs']:
        fn = os.path.basename(fn)
        ext = fn.rsplit('.', 1)[-1]
        final_name = os.path.join(CONFIG['movies_directory'], base_name + '.' + ext)
        if os.path.isfile(final_name):
            print("File already exists: %s" % final_name)
            sys.exit(1)

        # Move the files into place
        shutil.move(os.path.join(job_dir, fn), final_name)
        print("New file: %s (%s)" % (final_name, os.path.basename(fn)))

    # Remove the old directory
    shutil.rmtree(job_dir)

    # try to remove the empty category directories
    parent_dirname = os.path.dirname(job_dir)
    parent_basename = os.path.basename(parent_dirname)
    if parent_basename.lower() == category.lower():
        try:
            os.rmdir(parent_dirname)
            print("Removed empty directory: %s" % parent_dirname)
        except OSError:
            print("Skipped non-empty directory: %s" % parent_dirname)
            pass

if __name__ == '__main__':
    main(sys.argv[1], sys.argv[5])
<commit_msg>Fix bug where subtitles are ignored if they're not at the root of the job directory.<commit_after>#!/usr/bin/env python
from __future__ import print_function
import errno
import json
import os
import re
import shutil
import string
import sys

import requests

CONFIG_FILE = 'scripts.conf'
EXTENSIONS = ['avi', 'm4v', 'mkv', 'mp4']
SUB_EXTENSIONS = ['idx', 'sub', 'srt']
PATTERN = re.compile('^(.*)(\d{4})\.(.*)', re.I)
REPLACEMENTS = {
    ':': '',
}
ENDPOINT = 'https://api.themoviedb.org/3/search/movie'

def get_config(cfgfile):
    if not cfgfile.startswith('/'):
        cfgfile = os.path.join(os.path.dirname(sys.argv[0]), '..', cfgfile)

    config = {}
    with open(cfgfile) as fh:
        for line in fh:
            k, v = re.split('\s+=\s+', line, 1)
            config[k] = v.strip()
    return config

CONFIG = get_config(CONFIG_FILE)

def get_matching_files(dirname):
    # Loop through the files in ``dirname``, trying to find one which
    # matches PATTERN. Also build a list of files and their sizes.
    #
    # If no files match PATTERN but we do find one or more files with
    # a suitable extension we assume the largest file is the one we
    # want. We take this file and use the ``job_dir`` as the base for
    # the new file name.
    files = {
        'match': (),
        'by_size': [],
        'subs': [],
    }
    for root, _, fns in os.walk(dirname):
        for fn in fns:
            ext = fn.rsplit('.', 1)[-1].lower()
            if ext in SUB_EXTENSIONS:
                files['subs'].append(os.path.join(root, fn))

            if ext not in EXTENSIONS:
                continue

            abs = os.path.abspath(os.path.join(root, fn))
            files['by_size'].append((fn, os.path.getsize(abs)))
            match = PATTERN.match(fn)
            if match:
                files['match'] = (abs, match)
                break

    files['by_size'] = sorted(files['by_size'], key=lambda x: x[1])

    return files

def get_canonical_name(name, year):
    """
    Translates a dotted filename into the real name with spaces while
    preserving acronyms.
    """
    if not name.endswith('.'):
        name += '.'

    regex = '((?:[A-Z]\.)+)'
    acronyms = re.findall(regex, name)
    parts = re.split(regex, name)

    if not parts:
        return name

    canonical_name = []
    for part in parts:
        if part in acronyms:
            part = re.sub('\.$', '', part)
        else:
            part = part.replace('.', ' ')
        canonical_name.append(part.strip())
    canonical_name = ' '.join(canonical_name)

    # look up TMDb for the show name
    canonical_name = get_tmdb_name(canonical_name, year)
    if not canonical_name:
        return None

    # replace trailing dots
    canonical_name = re.sub('\.+$', '', canonical_name)

    # replace unsuitable characters
    for k, v in REPLACEMENTS.items():
        canonical_name = canonical_name.replace(k, v)

    return canonical_name

def get_tmdb_name(name, year):
    headers = {'Accept': 'application/json'}
    params = {
        'api_key': CONFIG['TMDB_API_KEY'],
        'query': name,
    }
    response = requests.get(ENDPOINT, params=params, headers=headers)
    results = json.loads(response.content.decode('utf-8'))['results']
    if not results:
        return None

    for r in results:
        if r['release_date'].split('-')[0] == year:
            return r['original_title']

def mkdirp(path, mode=0755):
    try:
        os.makedirs(path, mode=mode)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def main(job_dir, category):
    files = get_matching_files(job_dir)

    if files['match']:
        fn, match = files['match']
    else:
        # no files at all? just exit
        if not files['by_size']:
            print("No suitable files found")
            sys.exit()

        # found a potential match
        fn = files['by_size'][0][0]
        fn = os.path.join(job_dir, fn)
        match = PATTERN.match(os.path.basename(job_dir))
        if not match:
            print("No suitable files found")
            sys.exit()

    dotted_name, year = match.groups()[:2]
    canonical_name = get_canonical_name(dotted_name, year)

    # bail out if TMBd doesn't find anything
    if not canonical_name:
        print("Could not find files in TMDb for %s %s" % (dotted_name, year))
        sys.exit()

    dotted_name = canonical_name.replace(' ', '.')

    base_name = "%s.%s" % (dotted_name, year)
    for fn in [fn] + files['subs']:
        fn = os.path.basename(fn)
        ext = fn.rsplit('.', 1)[-1]
        final_name = os.path.join(CONFIG['movies_directory'], base_name + '.' + ext)
        if os.path.isfile(final_name):
            print("File already exists: %s" % final_name)
            sys.exit(1)

        # Move the files into place
        shutil.move(os.path.join(job_dir, fn), final_name)
        print("New file: %s (%s)" % (final_name, os.path.basename(fn)))

    # Remove the old directory
    shutil.rmtree(job_dir)

    # try to remove the empty category directories
    parent_dirname = os.path.dirname(job_dir)
    parent_basename = os.path.basename(parent_dirname)
    if parent_basename.lower() == category.lower():
        try:
            os.rmdir(parent_dirname)
            print("Removed empty directory: %s" % parent_dirname)
        except OSError:
            print("Skipped non-empty directory: %s" % parent_dirname)
            pass

if __name__ == '__main__':
    main(sys.argv[1], sys.argv[5])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Configure ``portage(5)``
'''

# Import python libs
import os
import shutil

# Import salt libs
import salt.utils

# Import third party libs
try:
    import portage
    HAS_PORTAGE = True
except ImportError:
    HAS_PORTAGE = False
    import sys
    if os.path.isdir('/usr/lib/portage/pym'):
        try:
            # In a virtualenv, the portage python path needs to be manually
            # added
            sys.path.insert(0, '/usr/lib/portage/pym')
            import portage
            HAS_PORTAGE = True
        except ImportError:
            pass

BASE_PATH = '/etc/portage/package.{0}'
SUPPORTED_CONFS = ('accept_keywords', 'env', 'license', 'mask', 'properties',
                   'unmask', 'use')

def __virtual__():
    '''
    Confirm this module is on a Gentoo based system.
    '''
    if HAS_PORTAGE and __grains__['os'] == 'Gentoo':
        return 'portage_config'
    return False

def _porttree():
    return portage.db[portage.root]['porttree']

def _p_to_cp(p):
    '''
    Convert a package name or a DEPEND atom to category/package format.
    Raises an exception if program name is ambiguous.
    '''
    ret = _porttree().dbapi.xmatch("match-all", p)
    if ret:
        return portage.cpv_getkey(ret[0])
    return None

def enforce_nice_config():
    '''
    Enforce a nice tree structure for /etc/portage/package.* configuration
    files.

    .. seealso::
       :py:func:`salt.modules.ebuild.ex_mod_init`
         for information on automatically running this when pkg is used.

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.enforce_nice_config
    '''
    _convert_all_package_confs_to_dir()
    _order_all_package_confs()

def _convert_all_package_confs_to_dir():
    '''
    Convert all /etc/portage/package.* configuration files to directories.
    '''
    for conf_file in SUPPORTED_CONFS:
        _package_conf_file_to_dir(conf_file)

def _order_all_package_confs():
    '''
    Place all entries in /etc/portage/package.* config dirs in the correct
    file.
    '''
    for conf_file in SUPPORTED_CONFS:
        _package_conf_ordering(conf_file)
    _unify_keywords()

def _unify_keywords():
    '''
    Merge /etc/portage/package.keywords and
    /etc/portage/package.accept_keywords.
    '''
    old_path = BASE_PATH.format('keywords')
    if os.path.exists(old_path):
        if os.path.isdir(old_path):
            for triplet in os.walk(old_path):
                for file_name in triplet[2]:
                    file_path = '{0}/{1}'.format(triplet[0], file_name)
                    with salt.utils.fopen(file_path) as fh_:
                        for line in fh_:
                            line = line.strip()
                            if line and not line.startswith('#'):
                                append_to_package_conf(
                                    'accept_keywords', string=line)
            shutil.rmtree(old_path)
        else:
            with salt.utils.fopen(old_path) as fh_:
                for line in fh_:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        append_to_package_conf('accept_keywords', string=line)
            os.remove(old_path)

def _package_conf_file_to_dir(file_name):
    '''
    Convert a config file to a config directory.
    '''
    if file_name in SUPPORTED_CONFS:
        path = BASE_PATH.format(file_name)
        if os.path.exists(path):
            if os.path.isdir(path):
                return False
            else:
                os.rename(path, path + '.tmpbak')
                os.mkdir(path, 0755)
                with salt.utils.fopen(path + '.tmpbak') as fh_:
                    for line in fh_:
                        line = line.strip()
                        if line and not line.startswith('#'):
                            append_to_package_conf(file_name, string=line)
                os.remove(path + '.tmpbak')
                return True
        else:
            os.mkdir(path, 0755)
            return True

def _package_conf_ordering(conf, clean=True, keep_backup=False):
    '''
    Move entries in the correct file.
    '''
    if conf in SUPPORTED_CONFS:
        rearrange = []
        path = BASE_PATH.format(conf)

        backup_files = []

        for triplet in os.walk(path):
            for file_name in triplet[2]:
                file_path = '{0}/{1}'.format(triplet[0], file_name)
                cp = triplet[0][len(path) + 1:] + '/' + file_name

                shutil.copy(file_path, file_path + '.bak')
                backup_files.append(file_path + '.bak')

                if cp[0] == '/' or cp.split('/') > 2:
                    rearrange.extend(list(salt.utils.fopen(file_path)))
                    os.remove(file_path)
                else:
                    new_contents = ''
                    with salt.utils.fopen(file_path, 'r+') as file_handler:
                        for line in file_handler:
                            try:
                                atom = line.strip().split()[0]
                            except IndexError:
                                new_contents += line
                            else:
                                if atom[0] == '#' or \
                                        portage.dep_getkey(atom) == cp:
                                    new_contents += line
                                else:
                                    rearrange.append(line.strip())
                        if len(new_contents) != 0:
                            file_handler.seek(0)
                            file_handler.truncate(len(new_contents))
                            file_handler.write(new_contents)

                    if len(new_contents) == 0:
                        os.remove(file_path)

        for line in rearrange:
            append_to_package_conf(conf, string=line)

        if not keep_backup:
            for bfile in backup_files:
                try:
                    os.remove(bfile)
                except OSError:
                    pass

        if clean:
            for triplet in os.walk(path):
                if len(triplet[1]) == 0 and len(triplet[2]) == 0 and \
                        triplet[0] != path:
                    shutil.rmtree(triplet[0])

def _merge_flags(*args):
    '''
    Merges multiple lists of flags removing duplicates and resolving conflicts
    giving priority to lasts lists.
    '''
    tmp = portage.flatten(args)
    flags = {}
    for flag in tmp:
        if flag[0] == '-':
            flags[flag[1:]] = False
        else:
            flags[flag] = True
    tmp = []
    for k, v in flags.iteritems():
        if v:
            tmp.append(k)
        else:
            tmp.append('-' + k)

    # Next sort is just aesthetic, can be commented for a small performance
    # boost
    tmp.sort(cmp=lambda x, y: cmp(x.lstrip('-'), y.lstrip('-')))
    return tmp

def append_to_package_conf(conf, atom='', flags=None, string='', overwrite=False):
    '''
    Append a string or a list of flags for a given package or DEPEND atom to a
    given configuration file.

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.append_to_package_conf use string="app-admin/salt ldap -libvirt"
        salt '*' portage_config.append_to_package_conf use atom="> = app-admin/salt-0.14.1" flags="['ldap', '-libvirt']"
    '''
    if flags is None:
        flags = []
    if conf in SUPPORTED_CONFS:
        if not string:
            if '/' not in atom:
                atom = _p_to_cp(atom)
                if not atom:
                    return
            string = '{0} {1}'.format(atom, ' '.join(flags))
            new_flags = list(flags)
        else:
            atom = string.strip().split()[0]
            new_flags = portage.dep.strip_empty(string.strip().split(' '))[1:]
            if '/' not in atom:
                atom = _p_to_cp(atom)
                string = '{0} {1}'.format(atom, ' '.join(new_flags))
                if not atom:
                    return

        to_delete_if_empty = []
        if conf == 'accept_keywords':
            if '-~ARCH' in new_flags:
                new_flags.remove('-~ARCH')
                to_delete_if_empty.append(atom)

            if '~ARCH' in new_flags:
                new_flags.remove('~ARCH')
                append_to_package_conf(conf, string=atom, overwrite=overwrite)
                if not new_flags:
                    return

        # Next sort is just aesthetic, can be commented for a small performance
        # boost
        new_flags.sort(cmp=lambda x, y: cmp(x.lstrip('-'), y.lstrip('-')))

        package_file = _p_to_cp(atom)
        if not package_file:
            return

        psplit = package_file.split('/')
        if len(psplit) == 2:
            pdir = BASE_PATH.format(conf) + '/' + psplit[0]
            if not os.path.exists(pdir):
                os.mkdir(pdir, 0755)

        complete_file_path = BASE_PATH.format(conf) + '/' + package_file

        try:
            shutil.copy(complete_file_path, complete_file_path + '.bak')
        except IOError:
            pass

        try:
            file_handler = salt.utils.fopen(complete_file_path, 'r+')
        except IOError:
            file_handler = salt.utils.fopen(complete_file_path, 'w+')

        new_contents = ''
        added = False

        for l in file_handler:
            l_strip = l.strip()
            if l_strip == '':
                new_contents += '\n'
            elif l_strip[0] == '#':
                new_contents += l
            elif l_strip.split()[0] == atom:
                if l_strip in to_delete_if_empty:
                    continue
                if overwrite:
                    new_contents += string.strip() + '\n'
                    added = True
                else:
                    old_flags = portage.dep.strip_empty(l_strip.split(' '))[1:]
                    if conf == 'accept_keywords':
                        if not old_flags:
                            new_contents += l
                            if not new_flags:
                                added = True
                            continue
                        elif not new_flags:
                            continue
                    merged_flags = _merge_flags(old_flags, new_flags)
                    if merged_flags:
                        new_contents += '{0} {1}\n'.format(
                            atom, ' '.join(merged_flags))
                    else:
                        new_contents += '{0}\n'.format(atom)
                    added = True
            else:
                new_contents += l
        if not added:
            new_contents += string.strip() + '\n'
        file_handler.seek(0)
        file_handler.truncate(len(new_contents))
        file_handler.write(new_contents)
        file_handler.close()
        try:
            os.remove(complete_file_path + '.bak')
        except OSError:
            pass

def append_use_flags(atom, uses=None, overwrite=False):
    '''
    Append a list of use flags for a given package or DEPEND atom

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.append_use_flags "app-admin/salt[ldap, -libvirt]"
        salt '*' portage_config.append_use_flags ">=app-admin/salt-0.14.1" "['ldap', '-libvirt']"
    '''
    if not uses:
        uses = portage.dep.dep_getusedeps(atom)
    if len(uses) == 0:
        return
    atom = atom[:atom.rfind('[')]
    append_to_package_conf('use', atom=atom, flags=uses, overwrite=overwrite)

def get_flags_from_package_conf(conf, atom):
    '''
    Get flags for a given package or DEPEND atom.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.get_flags_from_package_conf license salt
    '''
    if conf in SUPPORTED_CONFS:
        package_file = '{0}/{1}'.format(BASE_PATH.format(conf), _p_to_cp(atom))
        if '/' not in atom:
            atom = _p_to_cp(atom)
        try:
            match_list = set(_porttree().dbapi.xmatch("match-all", atom))
        except AttributeError:
            return []
        flags = []
        try:
            file_handler = salt.utils.fopen(package_file)
        except IOError:
            return []
        else:
            for line in file_handler:
                line = line.strip()
                line_package = line.split()[0]
                line_list = _porttree().dbapi.xmatch("match-all", line_package)
                if match_list.issubset(line_list):
                    f_tmp = portage.dep.strip_empty(line.strip().split()[1:])
                    if f_tmp:
                        flags.extend(f_tmp)
                    else:
                        flags.append('~ARCH')
            return _merge_flags(flags)

def has_flag(conf, atom, flag):
    '''
    Verify if the given package or DEPEND atom has the given flag.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.has_flag license salt Apache-2.0
    '''
    if flag in get_flags_from_package_conf(conf, atom):
        return True
    return False

def get_missing_flags(conf, atom, flags):
    '''
    Find out which of the given flags are currently not set.
    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.get_missing_flags use salt "['ldap', '-libvirt', 'openssl']"
    '''
    new_flags = []
    for flag in flags:
        if not has_flag(conf, atom, flag):
            new_flags.append(flag)
    return new_flags

def has_use(atom, use):
    '''
    Verify if the given package or DEPEND atom has the given use flag.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.has_use salt libvirt
    '''
    return has_flag('use', atom, use)

def is_present(conf, atom):
    '''
    Tell if a given package or DEPEND atom is present in the configuration
    files tree.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.is_present unmask salt
    '''
    if conf in SUPPORTED_CONFS:
        package_file = '{0}/{1}'.format(BASE_PATH.format(conf), _p_to_cp(atom))
        match_list = set(_porttree().dbapi.xmatch("match-all", atom))
        try:
            file_handler = salt.utils.fopen(package_file)
        except IOError:
            return False
        else:
            for line in file_handler:
                line = line.strip()
                line_package = line.split()[0]
                line_list = _porttree().dbapi.xmatch("match-all", line_package)
                if match_list.issubset(line_list):
                    return True
            return False
<commit_msg>py3 portage config<commit_after># -*- coding: utf-8 -*-
'''
Configure ``portage(5)``
'''

from __future__ import absolute_import

# Import python libs
import os
import shutil

# Import salt libs
import salt.utils

# Import third party libs
try:
    import portage
    HAS_PORTAGE = True
except ImportError:
    HAS_PORTAGE = False
    import sys
    if os.path.isdir('/usr/lib/portage/pym'):
        try:
            # In a virtualenv, the portage python path needs to be manually
            # added
            sys.path.insert(0, '/usr/lib/portage/pym')
            import portage
            HAS_PORTAGE = True
        except ImportError:
            pass

BASE_PATH = '/etc/portage/package.{0}'
SUPPORTED_CONFS = ('accept_keywords', 'env', 'license', 'mask', 'properties',
                   'unmask', 'use')

def __virtual__():
    '''
    Confirm this module is on a Gentoo based system.
    '''
    if HAS_PORTAGE and __grains__['os'] == 'Gentoo':
        return 'portage_config'
    return False

def _porttree():
    return portage.db[portage.root]['porttree']

def _p_to_cp(p):
    '''
    Convert a package name or a DEPEND atom to category/package format.
    Raises an exception if program name is ambiguous.
    '''
    ret = _porttree().dbapi.xmatch("match-all", p)
    if ret:
        return portage.cpv_getkey(ret[0])
    return None

def enforce_nice_config():
    '''
    Enforce a nice tree structure for /etc/portage/package.* configuration
    files.

    .. seealso::
       :py:func:`salt.modules.ebuild.ex_mod_init`
         for information on automatically running this when pkg is used.

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.enforce_nice_config
    '''
    _convert_all_package_confs_to_dir()
    _order_all_package_confs()

def _convert_all_package_confs_to_dir():
    '''
    Convert all /etc/portage/package.* configuration files to directories.
    '''
    for conf_file in SUPPORTED_CONFS:
        _package_conf_file_to_dir(conf_file)

def _order_all_package_confs():
    '''
    Place all entries in /etc/portage/package.* config dirs in the correct
    file.
    '''
    for conf_file in SUPPORTED_CONFS:
        _package_conf_ordering(conf_file)
    _unify_keywords()

def _unify_keywords():
    '''
    Merge /etc/portage/package.keywords and
    /etc/portage/package.accept_keywords.
    '''
    old_path = BASE_PATH.format('keywords')
    if os.path.exists(old_path):
        if os.path.isdir(old_path):
            for triplet in os.walk(old_path):
                for file_name in triplet[2]:
                    file_path = '{0}/{1}'.format(triplet[0], file_name)
                    with salt.utils.fopen(file_path) as fh_:
                        for line in fh_:
                            line = line.strip()
                            if line and not line.startswith('#'):
                                append_to_package_conf(
                                    'accept_keywords', string=line)
            shutil.rmtree(old_path)
        else:
            with salt.utils.fopen(old_path) as fh_:
                for line in fh_:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        append_to_package_conf('accept_keywords', string=line)
            os.remove(old_path)

def _package_conf_file_to_dir(file_name):
    '''
    Convert a config file to a config directory.
    '''
    if file_name in SUPPORTED_CONFS:
        path = BASE_PATH.format(file_name)
        if os.path.exists(path):
            if os.path.isdir(path):
                return False
            else:
                os.rename(path, path + '.tmpbak')
                os.mkdir(path, 0o755)
                with salt.utils.fopen(path + '.tmpbak') as fh_:
                    for line in fh_:
                        line = line.strip()
                        if line and not line.startswith('#'):
                            append_to_package_conf(file_name, string=line)
                os.remove(path + '.tmpbak')
                return True
        else:
            os.mkdir(path, 0o755)
            return True

def _package_conf_ordering(conf, clean=True, keep_backup=False):
    '''
    Move entries in the correct file.
    '''
    if conf in SUPPORTED_CONFS:
        rearrange = []
        path = BASE_PATH.format(conf)

        backup_files = []

        for triplet in os.walk(path):
            for file_name in triplet[2]:
                file_path = '{0}/{1}'.format(triplet[0], file_name)
                cp = triplet[0][len(path) + 1:] + '/' + file_name

                shutil.copy(file_path, file_path + '.bak')
                backup_files.append(file_path + '.bak')

                if cp[0] == '/' or cp.split('/') > 2:
                    rearrange.extend(list(salt.utils.fopen(file_path)))
                    os.remove(file_path)
                else:
                    new_contents = ''
                    with salt.utils.fopen(file_path, 'r+') as file_handler:
                        for line in file_handler:
                            try:
                                atom = line.strip().split()[0]
                            except IndexError:
                                new_contents += line
                            else:
                                if atom[0] == '#' or \
                                        portage.dep_getkey(atom) == cp:
                                    new_contents += line
                                else:
                                    rearrange.append(line.strip())
                        if len(new_contents) != 0:
                            file_handler.seek(0)
                            file_handler.truncate(len(new_contents))
                            file_handler.write(new_contents)

                    if len(new_contents) == 0:
                        os.remove(file_path)

        for line in rearrange:
            append_to_package_conf(conf, string=line)

        if not keep_backup:
            for bfile in backup_files:
                try:
                    os.remove(bfile)
                except OSError:
                    pass

        if clean:
            for triplet in os.walk(path):
                if len(triplet[1]) == 0 and len(triplet[2]) == 0 and \
                        triplet[0] != path:
                    shutil.rmtree(triplet[0])

def _merge_flags(*args):
    '''
    Merges multiple lists of flags removing duplicates and resolving conflicts
    giving priority to lasts lists.
    '''
    tmp = portage.flatten(args)
    flags = {}
    for flag in tmp:
        if flag[0] == '-':
            flags[flag[1:]] = False
        else:
            flags[flag] = True
    tmp = []
    for k, v in flags.items():
        if v:
            tmp.append(k)
        else:
            tmp.append('-' + k)

    # Next sort is just aesthetic, can be commented for a small performance
    # boost
    tmp.sort(cmp=lambda x, y: cmp(x.lstrip('-'), y.lstrip('-')))
    return tmp

def append_to_package_conf(conf, atom='', flags=None, string='', overwrite=False):
    '''
    Append a string or a list of flags for a given package or DEPEND atom to a
    given configuration file.

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.append_to_package_conf use string="app-admin/salt ldap -libvirt"
        salt '*' portage_config.append_to_package_conf use atom="> = app-admin/salt-0.14.1" flags="['ldap', '-libvirt']"
    '''
    if flags is None:
        flags = []
    if conf in SUPPORTED_CONFS:
        if not string:
            if '/' not in atom:
                atom = _p_to_cp(atom)
                if not atom:
                    return
            string = '{0} {1}'.format(atom, ' '.join(flags))
            new_flags = list(flags)
        else:
            atom = string.strip().split()[0]
            new_flags = portage.dep.strip_empty(string.strip().split(' '))[1:]
            if '/' not in atom:
                atom = _p_to_cp(atom)
                string = '{0} {1}'.format(atom, ' '.join(new_flags))
                if not atom:
                    return

        to_delete_if_empty = []
        if conf == 'accept_keywords':
            if '-~ARCH' in new_flags:
                new_flags.remove('-~ARCH')
                to_delete_if_empty.append(atom)

            if '~ARCH' in new_flags:
                new_flags.remove('~ARCH')
                append_to_package_conf(conf, string=atom, overwrite=overwrite)
                if not new_flags:
                    return

        # Next sort is just aesthetic, can be commented for a small performance
        # boost
        new_flags.sort(cmp=lambda x, y: cmp(x.lstrip('-'), y.lstrip('-')))

        package_file = _p_to_cp(atom)
        if not package_file:
            return

        psplit = package_file.split('/')
        if len(psplit) == 2:
            pdir = BASE_PATH.format(conf) + '/' + psplit[0]
            if not os.path.exists(pdir):
                os.mkdir(pdir, 0o755)

        complete_file_path = BASE_PATH.format(conf) + '/' + package_file

        try:
            shutil.copy(complete_file_path, complete_file_path + '.bak')
        except IOError:
            pass

        try:
            file_handler = salt.utils.fopen(complete_file_path, 'r+')
        except IOError:
            file_handler = salt.utils.fopen(complete_file_path, 'w+')

        new_contents = ''
        added = False

        for l in file_handler:
            l_strip = l.strip()
            if l_strip == '':
                new_contents += '\n'
            elif l_strip[0] == '#':
                new_contents += l
            elif l_strip.split()[0] == atom:
                if l_strip in to_delete_if_empty:
                    continue
                if overwrite:
                    new_contents += string.strip() + '\n'
                    added = True
                else:
                    old_flags = portage.dep.strip_empty(l_strip.split(' '))[1:]
                    if conf == 'accept_keywords':
                        if not old_flags:
                            new_contents += l
                            if not new_flags:
                                added = True
                            continue
                        elif not new_flags:
                            continue
                    merged_flags = _merge_flags(old_flags, new_flags)
                    if merged_flags:
                        new_contents += '{0} {1}\n'.format(
                            atom, ' '.join(merged_flags))
                    else:
                        new_contents += '{0}\n'.format(atom)
                    added = True
            else:
                new_contents += l
        if not added:
            new_contents += string.strip() + '\n'
        file_handler.seek(0)
        file_handler.truncate(len(new_contents))
        file_handler.write(new_contents)
        file_handler.close()
        try:
            os.remove(complete_file_path + '.bak')
        except OSError:
            pass

def append_use_flags(atom, uses=None, overwrite=False):
    '''
    Append a list of use flags for a given package or DEPEND atom

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.append_use_flags "app-admin/salt[ldap, -libvirt]"
        salt '*' portage_config.append_use_flags ">=app-admin/salt-0.14.1" "['ldap', '-libvirt']"
    '''
    if not uses:
        uses = portage.dep.dep_getusedeps(atom)
    if len(uses) == 0:
        return
    atom = atom[:atom.rfind('[')]
    append_to_package_conf('use', atom=atom, flags=uses, overwrite=overwrite)

def get_flags_from_package_conf(conf, atom):
    '''
    Get flags for a given package or DEPEND atom.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.get_flags_from_package_conf license salt
    '''
    if conf in SUPPORTED_CONFS:
        package_file = '{0}/{1}'.format(BASE_PATH.format(conf), _p_to_cp(atom))
        if '/' not in atom:
            atom = _p_to_cp(atom)
        try:
            match_list = set(_porttree().dbapi.xmatch("match-all", atom))
        except AttributeError:
            return []
        flags = []
        try:
            file_handler = salt.utils.fopen(package_file)
        except IOError:
            return []
        else:
            for line in file_handler:
                line = line.strip()
                line_package = line.split()[0]
                line_list = _porttree().dbapi.xmatch("match-all", line_package)
                if match_list.issubset(line_list):
                    f_tmp = portage.dep.strip_empty(line.strip().split()[1:])
                    if f_tmp:
                        flags.extend(f_tmp)
                    else:
                        flags.append('~ARCH')
            return _merge_flags(flags)

def has_flag(conf, atom, flag):
    '''
    Verify if the given package or DEPEND atom has the given flag.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.has_flag license salt Apache-2.0
    '''
    if flag in get_flags_from_package_conf(conf, atom):
        return True
    return False

def get_missing_flags(conf, atom, flags):
    '''
    Find out which of the given flags are currently not set.
    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.get_missing_flags use salt "['ldap', '-libvirt', 'openssl']"
    '''
    new_flags = []
    for flag in flags:
        if not has_flag(conf, atom, flag):
            new_flags.append(flag)
    return new_flags

def has_use(atom, use):
    '''
    Verify if the given package or DEPEND atom has the given use flag.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.has_use salt libvirt
    '''
    return has_flag('use', atom, use)

def is_present(conf, atom):
    '''
    Tell if a given package or DEPEND atom is present in the configuration
    files tree.
    Warning: This only works if the configuration files tree is in the correct
    format (the one enforced by enforce_nice_config)

    CLI Example:

    .. code-block:: bash

        salt '*' portage_config.is_present unmask salt
    '''
    if conf in SUPPORTED_CONFS:
        package_file = '{0}/{1}'.format(BASE_PATH.format(conf), _p_to_cp(atom))
        match_list = set(_porttree().dbapi.xmatch("match-all", atom))
        try:
            file_handler = salt.utils.fopen(package_file)
        except IOError:
            return False
        else:
            for line in file_handler:
                line = line.strip()
                line_package = line.split()[0]
                line_list = _porttree().dbapi.xmatch("match-all", line_package)
                if match_list.issubset(line_list):
                    return True
            return False
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Create virtualenv environments.

.. versionadded:: 0.17.0
'''

# Import python libs
from __future__ import absolute_import
import glob
import shutil
import logging
import os
import sys

# Import salt libs
import salt.utils
import salt.utils.files
import salt.utils.verify
from salt.exceptions import CommandExecutionError, SaltInvocationError
from salt.ext.six import string_types

KNOWN_BINARY_NAMES = frozenset([
    'virtualenv-{0}.{1}'.format(*sys.version_info[:2]),
    'virtualenv{0}'.format(sys.version_info[0]),
    'virtualenv',
])

log = logging.getLogger(__name__)

__opts__ = {
    'venv_bin': salt.utils.which_bin(KNOWN_BINARY_NAMES) or 'virtualenv'
}

__pillar__ = {}

# Define the module's virtual name
__virtualname__ = 'virtualenv'

def __virtual__():
    return __virtualname__

def create(path,
           venv_bin=None,
           system_site_packages=False,
           distribute=False,
           clear=False,
           python=None,
           extra_search_dir=None,
           never_download=None,
           prompt=None,
           pip=False,
           symlinks=None,
           upgrade=None,
           user=None,
           use_vt=False,
           saltenv='base'):
    '''
    Create a virtualenv

    path
        The path to the virtualenv to be created

    venv_bin
        The name (and optionally path) of the virtualenv command. This can also
        be set globally in the minion config file as ``virtualenv.venv_bin``.
        Defaults to ``virtualenv``.

    system_site_packages : False
        Passthrough argument given to virtualenv or pyvenv

    distribute : False
        Passthrough argument given to virtualenv

    pip : False
        Install pip after creating a virtual environment. Implies
        ``distribute=True``

    clear : False
        Passthrough argument given to virtualenv or pyvenv

    python : None (default)
        Passthrough argument given to virtualenv

    extra_search_dir : None (default)
        Passthrough argument given to virtualenv

    never_download : None (default)
        Passthrough argument given to virtualenv if True

    prompt : None (default)
        Passthrough argument given to virtualenv if not None

    symlinks : None
        Passthrough argument given to pyvenv if True

    upgrade : None
        Passthrough argument given to pyvenv if True

    user : None
        Set ownership for the virtualenv

    runas : None
        Set ownership for the virtualenv

        .. deprecated:: 2014.1.0
            ``user`` should be used instead

    use_vt : False
        Use VT terminal emulation (see output while installing)

        .. versionadded:: 2015.5.0

    saltenv : 'base'
        Specify a different environment. The default environment is ``base``.

        .. versionadded:: 2014.1.0

    .. note::
        The ``runas`` argument is deprecated as of 2014.1.0. ``user`` should be
        used instead.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.create /path/to/new/virtualenv
    '''
    if venv_bin is None:
        venv_bin = __opts__.get('venv_bin') or __pillar__.get('venv_bin')

    cmd = [venv_bin]

    if 'pyvenv' not in venv_bin:
        # ----- Stop the user if pyvenv only options are used --------------->
        # If any of the following values are not None, it means that the user
        # is actually passing a True or False value. Stop Him!
        if upgrade is not None:
            raise CommandExecutionError(
                'The `upgrade`(`--upgrade`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        elif symlinks is not None:
            raise CommandExecutionError(
                'The `symlinks`(`--symlinks`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        # <---- Stop the user if pyvenv only options are used ----------------

        # Virtualenv package
        try:
            import virtualenv
            version = getattr(virtualenv, '__version__',
                              virtualenv.virtualenv_version)
            virtualenv_version_info = tuple(
                [int(i) for i in version.split('rc')[0].split('.')]
            )
        except ImportError:
            # Unable to import?? Let's parse the version from the console
            version_cmd = [venv_bin, '--version']
            ret = __salt__['cmd.run_all'](
                    version_cmd, runas=user, python_shell=False
                )
            if ret['retcode'] > 0 or not ret['stdout'].strip():
                raise CommandExecutionError(
                    'Unable to get the virtualenv version output using \'{0}\'. '
                    'Returned data: {1}'.format(version_cmd, ret)
                )
            virtualenv_version_info = tuple(
                [int(i) for i in
                 ret['stdout'].strip().split('rc')[0].split('.')]
            )

        if distribute:
            if virtualenv_version_info >= (1, 10):
                log.info(
                    'The virtualenv \'--distribute\' option has been '
                    'deprecated in virtualenv(>=1.10), as such, the '
                    '\'distribute\' option to `virtualenv.create()` has '
                    'also been deprecated and it\'s not necessary anymore.'
                )
            else:
                cmd.append('--distribute')

        if python is not None and python.strip() != '':
            if not salt.utils.which(python):
                raise CommandExecutionError(
                    'Cannot find requested python ({0}).'.format(python)
                )
            cmd.append('--python={0}'.format(python))
        if extra_search_dir is not None:
            if isinstance(extra_search_dir, string_types) and \
                    extra_search_dir.strip() != '':
                extra_search_dir = [
                    e.strip() for e in extra_search_dir.split(',')
                ]
            for entry in extra_search_dir:
                cmd.append('--extra-search-dir={0}'.format(entry))
        if never_download is True:
            if virtualenv_version_info >= (1, 10) and virtualenv_version_info < (14, 0, 0):
                log.info(
                    '--never-download was deprecated in 1.10.0, but reimplemented in 14.0.0. '
                    'If this feature is needed, please install a supported virtualenv version.'
                )
            else:
                cmd.append('--never-download')
        if prompt is not None and prompt.strip() != '':
            cmd.append('--prompt=\'{0}\''.format(prompt))
    else:
        # venv module from the Python >= 3.3 standard library

        # ----- Stop the user if virtualenv only options are being used ----->
        # If any of the following values are not None, it means that the user
        # is actually passing a True or False value. Stop Him!
        if python is not None and python.strip() != '':
            raise CommandExecutionError(
                'The `python`(`--python`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        elif extra_search_dir is not None and extra_search_dir.strip() != '':
            raise CommandExecutionError(
                'The `extra_search_dir`(`--extra-search-dir`) option is not '
                'supported by \'{0}\''.format(venv_bin)
            )
        elif never_download is not None:
            raise CommandExecutionError(
                'The `never_download`(`--never-download`) option is not '
                'supported by \'{0}\''.format(venv_bin)
            )
        elif prompt is not None and prompt.strip() != '':
            raise CommandExecutionError(
                'The `prompt`(`--prompt`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        # <---- Stop the user if virtualenv only options are being used ------

        if upgrade is True:
            cmd.append('--upgrade')
        if symlinks is True:
            cmd.append('--symlinks')

    # Common options to virtualenv and pyvenv
    if clear is True:
        cmd.append('--clear')
    if system_site_packages is True:
        cmd.append('--system-site-packages')

    # Finally the virtualenv path
    cmd.append(path)

    # Let's create the virtualenv
    ret = __salt__['cmd.run_all'](cmd, runas=user, python_shell=False)
    if ret['retcode'] != 0:
        # Something went wrong. Let's bail out now!
        return ret

    # Check if distribute and pip are already installed
    if salt.utils.is_windows():
        venv_python = os.path.join(path, 'Scripts', 'python.exe')
        venv_pip = os.path.join(path, 'Scripts', 'pip.exe')
        venv_setuptools = os.path.join(path, 'Scripts', 'easy_install.exe')
    else:
        venv_python = os.path.join(path, 'bin', 'python')
        venv_pip = os.path.join(path, 'bin', 'pip')
        venv_setuptools = os.path.join(path, 'bin', 'easy_install')

    # Install setuptools
    if (pip or distribute) and not os.path.exists(venv_setuptools):
        _install_script(
            'https://bitbucket.org/pypa/setuptools/raw/default/ez_setup.py',
            path, venv_python, user, saltenv=saltenv, use_vt=use_vt
        )

        # clear up the distribute archive which gets downloaded
        for fpath in glob.glob(os.path.join(path, 'distribute-*.tar.gz*')):
            os.unlink(fpath)

    if ret['retcode'] != 0:
        # Something went wrong. Let's bail out now!
        return ret

    # Install pip
    if pip and not os.path.exists(venv_pip):
        _ret = _install_script(
            'https://bootstrap.pypa.io/get-pip.py',
            path, venv_python, user, saltenv=saltenv, use_vt=use_vt
        )
        # Let's update the return dictionary with the details from the pip
        # installation
        ret.update(
            retcode=_ret['retcode'],
            stdout='{0}\n{1}'.format(ret['stdout'], _ret['stdout']).strip(),
            stderr='{0}\n{1}'.format(ret['stderr'], _ret['stderr']).strip(),
        )

    return ret

def get_site_packages(venv):
    '''
    Return the path to the site-packages directory of a virtualenv

    venv
        Path to the virtualenv.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_site_packages /path/to/my/venv
    '''
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'from distutils import sysconfig; '
            'print sysconfig.get_python_lib()'
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_distribution_path(venv, distribution):
    '''
    Return the path to a distribution installed inside a virtualenv

    .. versionadded:: 2016.3.0

    venv
        Path to the virtualenv.
    distribution
        Name of the distribution. Note, all non-alphanumeric characters
        will be converted to dashes.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_distribution_path /path/to/my/venv my_distribution
    '''
    _verify_safe_py_code(distribution)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.get_distribution('{0}').location)".format(
                distribution
            )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_resource_path(venv,
                      package=None,
                      resource=None):
    '''
    Return the path to a package resource installed inside a virtualenv

    .. versionadded:: 2015.5.0

    venv
        Path to the virtualenv

    package
        Name of the package in which the resource resides

        .. versionadded:: 2016.3.0

    resource
        Name of the resource of which the path is to be returned

        .. versionadded:: 2016.3.0

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_resource_path /path/to/my/venv my_package my/resource.xml
    '''
    _verify_safe_py_code(package, resource)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.resource_filename('{0}', '{1}'))".format(
                package,
                resource
        )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_resource_content(venv,
                         package=None,
                         resource=None):
    '''
    Return the content of a package resource installed inside a virtualenv

    .. versionadded:: 2015.5.0

    venv
        Path to the virtualenv

    package
        Name of the package in which the resource resides

        .. versionadded:: 2016.3.0

    resource
        Name of the resource of which the content is to be returned

        .. versionadded:: 2016.3.0

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_resource_content /path/to/my/venv my_package my/resource.xml
    '''
    _verify_safe_py_code(package, resource)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.resource_string('{0}', '{1}'))".format(
                package,
                resource
            )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def _install_script(source, cwd, python, user, saltenv='base', use_vt=False):
    if not salt.utils.is_windows():
        tmppath = salt.utils.files.mkstemp(dir=cwd)
    else:
        tmppath = __salt__['cp.cache_file'](source, saltenv)

    if not salt.utils.is_windows():
        fn_ = __salt__['cp.cache_file'](source, saltenv)
        shutil.copyfile(fn_, tmppath)
        os.chmod(tmppath, 0o500)
        os.chown(tmppath, __salt__['file.user_to_uid'](user), -1)
    try:
        return __salt__['cmd.run_all'](
            [python, tmppath],
            runas=user,
            cwd=cwd,
            env={'VIRTUAL_ENV': cwd},
            use_vt=use_vt,
            python_shell=False,
        )
    finally:
        os.remove(tmppath)

def _verify_safe_py_code(*args):
    for arg in args:
        if not salt.utils.verify.safe_py_code(arg):
            raise SaltInvocationError(
                'Unsafe python code detected in \'{0}\''.format(arg)
            )

def _verify_virtualenv(venv_path):
    bin_path = os.path.join(venv_path, 'bin/python')
    if not os.path.exists(bin_path):
        raise CommandExecutionError(
            'Path \'{0}\' does not appear to be a virtualenv: bin/python not found.'.format(venv_path)
        )
    return bin_path
<commit_msg>add missing parenthis to keep integration with python3<commit_after># -*- coding: utf-8 -*-
'''
Create virtualenv environments.

.. versionadded:: 0.17.0
'''

# Import python libs
from __future__ import absolute_import
import glob
import shutil
import logging
import os
import sys

# Import salt libs
import salt.utils
import salt.utils.files
import salt.utils.verify
from salt.exceptions import CommandExecutionError, SaltInvocationError
from salt.ext.six import string_types

KNOWN_BINARY_NAMES = frozenset([
    'virtualenv-{0}.{1}'.format(*sys.version_info[:2]),
    'virtualenv{0}'.format(sys.version_info[0]),
    'virtualenv',
])

log = logging.getLogger(__name__)

__opts__ = {
    'venv_bin': salt.utils.which_bin(KNOWN_BINARY_NAMES) or 'virtualenv'
}

__pillar__ = {}

# Define the module's virtual name
__virtualname__ = 'virtualenv'

def __virtual__():
    return __virtualname__

def create(path,
           venv_bin=None,
           system_site_packages=False,
           distribute=False,
           clear=False,
           python=None,
           extra_search_dir=None,
           never_download=None,
           prompt=None,
           pip=False,
           symlinks=None,
           upgrade=None,
           user=None,
           use_vt=False,
           saltenv='base'):
    '''
    Create a virtualenv

    path
        The path to the virtualenv to be created

    venv_bin
        The name (and optionally path) of the virtualenv command. This can also
        be set globally in the minion config file as ``virtualenv.venv_bin``.
        Defaults to ``virtualenv``.

    system_site_packages : False
        Passthrough argument given to virtualenv or pyvenv

    distribute : False
        Passthrough argument given to virtualenv

    pip : False
        Install pip after creating a virtual environment. Implies
        ``distribute=True``

    clear : False
        Passthrough argument given to virtualenv or pyvenv

    python : None (default)
        Passthrough argument given to virtualenv

    extra_search_dir : None (default)
        Passthrough argument given to virtualenv

    never_download : None (default)
        Passthrough argument given to virtualenv if True

    prompt : None (default)
        Passthrough argument given to virtualenv if not None

    symlinks : None
        Passthrough argument given to pyvenv if True

    upgrade : None
        Passthrough argument given to pyvenv if True

    user : None
        Set ownership for the virtualenv

    runas : None
        Set ownership for the virtualenv

        .. deprecated:: 2014.1.0
            ``user`` should be used instead

    use_vt : False
        Use VT terminal emulation (see output while installing)

        .. versionadded:: 2015.5.0

    saltenv : 'base'
        Specify a different environment. The default environment is ``base``.

        .. versionadded:: 2014.1.0

    .. note::
        The ``runas`` argument is deprecated as of 2014.1.0. ``user`` should be
        used instead.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.create /path/to/new/virtualenv
    '''
    if venv_bin is None:
        venv_bin = __opts__.get('venv_bin') or __pillar__.get('venv_bin')

    cmd = [venv_bin]

    if 'pyvenv' not in venv_bin:
        # ----- Stop the user if pyvenv only options are used --------------->
        # If any of the following values are not None, it means that the user
        # is actually passing a True or False value. Stop Him!
        if upgrade is not None:
            raise CommandExecutionError(
                'The `upgrade`(`--upgrade`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        elif symlinks is not None:
            raise CommandExecutionError(
                'The `symlinks`(`--symlinks`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        # <---- Stop the user if pyvenv only options are used ----------------

        # Virtualenv package
        try:
            import virtualenv
            version = getattr(virtualenv, '__version__',
                              virtualenv.virtualenv_version)
            virtualenv_version_info = tuple(
                [int(i) for i in version.split('rc')[0].split('.')]
            )
        except ImportError:
            # Unable to import?? Let's parse the version from the console
            version_cmd = [venv_bin, '--version']
            ret = __salt__['cmd.run_all'](
                    version_cmd, runas=user, python_shell=False
                )
            if ret['retcode'] > 0 or not ret['stdout'].strip():
                raise CommandExecutionError(
                    'Unable to get the virtualenv version output using \'{0}\'. '
                    'Returned data: {1}'.format(version_cmd, ret)
                )
            virtualenv_version_info = tuple(
                [int(i) for i in
                 ret['stdout'].strip().split('rc')[0].split('.')]
            )

        if distribute:
            if virtualenv_version_info >= (1, 10):
                log.info(
                    'The virtualenv \'--distribute\' option has been '
                    'deprecated in virtualenv(>=1.10), as such, the '
                    '\'distribute\' option to `virtualenv.create()` has '
                    'also been deprecated and it\'s not necessary anymore.'
                )
            else:
                cmd.append('--distribute')

        if python is not None and python.strip() != '':
            if not salt.utils.which(python):
                raise CommandExecutionError(
                    'Cannot find requested python ({0}).'.format(python)
                )
            cmd.append('--python={0}'.format(python))
        if extra_search_dir is not None:
            if isinstance(extra_search_dir, string_types) and \
                    extra_search_dir.strip() != '':
                extra_search_dir = [
                    e.strip() for e in extra_search_dir.split(',')
                ]
            for entry in extra_search_dir:
                cmd.append('--extra-search-dir={0}'.format(entry))
        if never_download is True:
            if virtualenv_version_info >= (1, 10) and virtualenv_version_info < (14, 0, 0):
                log.info(
                    '--never-download was deprecated in 1.10.0, but reimplemented in 14.0.0. '
                    'If this feature is needed, please install a supported virtualenv version.'
                )
            else:
                cmd.append('--never-download')
        if prompt is not None and prompt.strip() != '':
            cmd.append('--prompt=\'{0}\''.format(prompt))
    else:
        # venv module from the Python >= 3.3 standard library

        # ----- Stop the user if virtualenv only options are being used ----->
        # If any of the following values are not None, it means that the user
        # is actually passing a True or False value. Stop Him!
        if python is not None and python.strip() != '':
            raise CommandExecutionError(
                'The `python`(`--python`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        elif extra_search_dir is not None and extra_search_dir.strip() != '':
            raise CommandExecutionError(
                'The `extra_search_dir`(`--extra-search-dir`) option is not '
                'supported by \'{0}\''.format(venv_bin)
            )
        elif never_download is not None:
            raise CommandExecutionError(
                'The `never_download`(`--never-download`) option is not '
                'supported by \'{0}\''.format(venv_bin)
            )
        elif prompt is not None and prompt.strip() != '':
            raise CommandExecutionError(
                'The `prompt`(`--prompt`) option is not supported '
                'by \'{0}\''.format(venv_bin)
            )
        # <---- Stop the user if virtualenv only options are being used ------

        if upgrade is True:
            cmd.append('--upgrade')
        if symlinks is True:
            cmd.append('--symlinks')

    # Common options to virtualenv and pyvenv
    if clear is True:
        cmd.append('--clear')
    if system_site_packages is True:
        cmd.append('--system-site-packages')

    # Finally the virtualenv path
    cmd.append(path)

    # Let's create the virtualenv
    ret = __salt__['cmd.run_all'](cmd, runas=user, python_shell=False)
    if ret['retcode'] != 0:
        # Something went wrong. Let's bail out now!
        return ret

    # Check if distribute and pip are already installed
    if salt.utils.is_windows():
        venv_python = os.path.join(path, 'Scripts', 'python.exe')
        venv_pip = os.path.join(path, 'Scripts', 'pip.exe')
        venv_setuptools = os.path.join(path, 'Scripts', 'easy_install.exe')
    else:
        venv_python = os.path.join(path, 'bin', 'python')
        venv_pip = os.path.join(path, 'bin', 'pip')
        venv_setuptools = os.path.join(path, 'bin', 'easy_install')

    # Install setuptools
    if (pip or distribute) and not os.path.exists(venv_setuptools):
        _install_script(
            'https://bitbucket.org/pypa/setuptools/raw/default/ez_setup.py',
            path, venv_python, user, saltenv=saltenv, use_vt=use_vt
        )

        # clear up the distribute archive which gets downloaded
        for fpath in glob.glob(os.path.join(path, 'distribute-*.tar.gz*')):
            os.unlink(fpath)

    if ret['retcode'] != 0:
        # Something went wrong. Let's bail out now!
        return ret

    # Install pip
    if pip and not os.path.exists(venv_pip):
        _ret = _install_script(
            'https://bootstrap.pypa.io/get-pip.py',
            path, venv_python, user, saltenv=saltenv, use_vt=use_vt
        )
        # Let's update the return dictionary with the details from the pip
        # installation
        ret.update(
            retcode=_ret['retcode'],
            stdout='{0}\n{1}'.format(ret['stdout'], _ret['stdout']).strip(),
            stderr='{0}\n{1}'.format(ret['stderr'], _ret['stderr']).strip(),
        )

    return ret

def get_site_packages(venv):
    '''
    Return the path to the site-packages directory of a virtualenv

    venv
        Path to the virtualenv.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_site_packages /path/to/my/venv
    '''
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'from distutils import sysconfig; '
            'print(sysconfig.get_python_lib())'
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_distribution_path(venv, distribution):
    '''
    Return the path to a distribution installed inside a virtualenv

    .. versionadded:: 2016.3.0

    venv
        Path to the virtualenv.
    distribution
        Name of the distribution. Note, all non-alphanumeric characters
        will be converted to dashes.

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_distribution_path /path/to/my/venv my_distribution
    '''
    _verify_safe_py_code(distribution)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.get_distribution('{0}').location)".format(
                distribution
            )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_resource_path(venv,
                      package=None,
                      resource=None):
    '''
    Return the path to a package resource installed inside a virtualenv

    .. versionadded:: 2015.5.0

    venv
        Path to the virtualenv

    package
        Name of the package in which the resource resides

        .. versionadded:: 2016.3.0

    resource
        Name of the resource of which the path is to be returned

        .. versionadded:: 2016.3.0

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_resource_path /path/to/my/venv my_package my/resource.xml
    '''
    _verify_safe_py_code(package, resource)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.resource_filename('{0}', '{1}'))".format(
                package,
                resource
        )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def get_resource_content(venv,
                         package=None,
                         resource=None):
    '''
    Return the content of a package resource installed inside a virtualenv

    .. versionadded:: 2015.5.0

    venv
        Path to the virtualenv

    package
        Name of the package in which the resource resides

        .. versionadded:: 2016.3.0

    resource
        Name of the resource of which the content is to be returned

        .. versionadded:: 2016.3.0

    CLI Example:

    .. code-block:: bash

        salt '*' virtualenv.get_resource_content /path/to/my/venv my_package my/resource.xml
    '''
    _verify_safe_py_code(package, resource)
    bin_path = _verify_virtualenv(venv)

    ret = __salt__['cmd.exec_code_all'](
        bin_path,
        'import pkg_resources; '
            "print(pkg_resources.resource_string('{0}', '{1}'))".format(
                package,
                resource
            )
    )

    if ret['retcode'] != 0:
        raise CommandExecutionError('{stdout}\n{stderr}'.format(**ret))

    return ret['stdout']

def _install_script(source, cwd, python, user, saltenv='base', use_vt=False):
    if not salt.utils.is_windows():
        tmppath = salt.utils.files.mkstemp(dir=cwd)
    else:
        tmppath = __salt__['cp.cache_file'](source, saltenv)

    if not salt.utils.is_windows():
        fn_ = __salt__['cp.cache_file'](source, saltenv)
        shutil.copyfile(fn_, tmppath)
        os.chmod(tmppath, 0o500)
        os.chown(tmppath, __salt__['file.user_to_uid'](user), -1)
    try:
        return __salt__['cmd.run_all'](
            [python, tmppath],
            runas=user,
            cwd=cwd,
            env={'VIRTUAL_ENV': cwd},
            use_vt=use_vt,
            python_shell=False,
        )
    finally:
        os.remove(tmppath)

def _verify_safe_py_code(*args):
    for arg in args:
        if not salt.utils.verify.safe_py_code(arg):
            raise SaltInvocationError(
                'Unsafe python code detected in \'{0}\''.format(arg)
            )

def _verify_virtualenv(venv_path):
    bin_path = os.path.join(venv_path, 'bin/python')
    if not os.path.exists(bin_path):
        raise CommandExecutionError(
            'Path \'{0}\' does not appear to be a virtualenv: bin/python not found.'.format(venv_path)
        )
    return bin_path
<|endoftext|>