<commit_before>from __future__ import unicode_literals

import os
import re
import io
import getpass
import posixpath
import logging
from subprocess import Popen, PIPE
from contextlib import contextmanager, closing

from dvc.config import Config
from dvc.scheme import Schemes

from dvc.utils.compat import urlparse, FileNotFoundError
from dvc.utils import fix_env, tmp_fname

from .pool import get_connection
from .base import RemoteBASE, RemoteCmdError

logger = logging.getLogger(__name__)

class RemoteHDFS(RemoteBASE):
    scheme = Schemes.HDFS
    REGEX = r"^hdfs://((?P<user>.*)@)?.*$"
    PARAM_CHECKSUM = "checksum"
    REQUIRES = {"pyarrow": "pyarrow"}

    def __init__(self, repo, config):
        super(RemoteHDFS, self).__init__(repo, config)
        self.path_info = None
        url = config.get(Config.SECTION_REMOTE_URL)
        if not url:
            return

        parsed = urlparse(url)

        user = (
            parsed.username
            or config.get(Config.SECTION_REMOTE_USER)
            or getpass.getuser()
        )

        self.path_info = self.path_cls.from_parts(
            scheme=self.scheme,
            host=parsed.hostname,
            user=user,
            port=parsed.port,
            path=parsed.path,
        )

    def hdfs(self, path_info):
        import pyarrow

        return get_connection(
            pyarrow.hdfs.connect,
            path_info.host,
            path_info.port,
            user=path_info.user,
        )

    def hadoop_fs(self, cmd, user=None):
        cmd = "hadoop fs -" + cmd
        if user:
            cmd = "HADOOP_USER_NAME={} ".format(user) + cmd

        # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
        # See https://github.com/iterative/dvc/issues/1197.
        close_fds = os.name != "nt"

        executable = os.getenv("SHELL") if os.name != "nt" else None
        p = Popen(
            cmd,
            shell=True,
            close_fds=close_fds,
            executable=executable,
            env=fix_env(os.environ),
            stdin=PIPE,
            stdout=PIPE,
            stderr=PIPE,
        )
        out, err = p.communicate()
        if p.returncode != 0:
            raise RemoteCmdError(self.scheme, cmd, p.returncode, err)
        return out.decode("utf-8")

    @staticmethod
    def _group(regex, s, gname):
        match = re.match(regex, s)
        assert match is not None
        return match.group(gname)

    def get_file_checksum(self, path_info):
        # NOTE: pyarrow doesn't support checksum, so we need to use hadoop
        regex = r".*\t.*\t(?P<checksum>.*)"
        stdout = self.hadoop_fs(
            "checksum {}".format(path_info.path), user=path_info.user
        )
        return self._group(regex, stdout, "checksum")

    def copy(self, from_info, to_info, **_kwargs):
        dname = posixpath.dirname(to_info.path)
        with self.hdfs(to_info) as hdfs:
            hdfs.mkdir(dname)
            # NOTE: this is how `hadoop fs -cp` works too: it copies through
            # your local machine.
            with hdfs.open(from_info.path, "rb") as from_fobj:
                with hdfs.open(to_info.path, "wb") as to_fobj:
                    to_fobj.upload(from_fobj)

    def remove(self, path_info):
        if path_info.scheme != "hdfs":
            raise NotImplementedError

        if self.exists(path_info):
            logger.debug("Removing {}".format(path_info.path))
            with self.hdfs(path_info) as hdfs:
                hdfs.rm(path_info.path)

    def exists(self, path_info):
        assert not isinstance(path_info, list)
        assert path_info.scheme == "hdfs"
        with self.hdfs(path_info) as hdfs:
            return hdfs.exists(path_info.path)

    def _upload(self, from_file, to_info, **_kwargs):
        with self.hdfs(to_info) as hdfs:
            hdfs.mkdir(posixpath.dirname(to_info.path))
            tmp_file = tmp_fname(to_info.path)
            with open(from_file, "rb") as fobj:
                hdfs.upload(tmp_file, fobj)
            hdfs.rename(tmp_file, to_info.path)

    def _download(self, from_info, to_file, **_kwargs):
        with self.hdfs(from_info) as hdfs:
            with open(to_file, "wb+") as fobj:
                hdfs.download(from_info.path, fobj)

    @contextmanager
    def open(self, path_info, mode="r", encoding=None):
        assert mode in {"r", "rt", "rb"}

        try:
            with self.hdfs(path_info) as hdfs, closing(
                hdfs.open(path_info.path, mode="rb")
            ) as fd:
                if mode == "rb":
                    yield fd
                else:
                    yield io.TextIOWrapper(fd, encoding=encoding)
        except IOError as e:
            # Empty .errno and not specific enough error class in pyarrow,
            # see https://issues.apache.org/jira/browse/ARROW-6248
            if "file does not exist" in str(e):
                raise FileNotFoundError(*e.args)
            raise

    def list_cache_paths(self):
        if not self.exists(self.path_info):
            return []

        files = []
        dirs = [self.path_info.path]

        with self.hdfs(self.path_info) as hdfs:
            while dirs:
                for entry in hdfs.ls(dirs.pop(), detail=True):
                    if entry["kind"] == "directory":
                        dirs.append(urlparse(entry["name"]).path)
                    elif entry["kind"] == "file":
                        files.append(urlparse(entry["name"]).path)

        return files
<commit_msg>hdfs: don't set default user (#2622)<commit_after>from __future__ import unicode_literals

import os
import re
import io
import posixpath
import logging
from subprocess import Popen, PIPE
from contextlib import contextmanager, closing

from dvc.config import Config
from dvc.scheme import Schemes

from dvc.utils.compat import urlparse, FileNotFoundError
from dvc.utils import fix_env, tmp_fname

from .pool import get_connection
from .base import RemoteBASE, RemoteCmdError

logger = logging.getLogger(__name__)

class RemoteHDFS(RemoteBASE):
    scheme = Schemes.HDFS
    REGEX = r"^hdfs://((?P<user>.*)@)?.*$"
    PARAM_CHECKSUM = "checksum"
    REQUIRES = {"pyarrow": "pyarrow"}

    def __init__(self, repo, config):
        super(RemoteHDFS, self).__init__(repo, config)
        self.path_info = None
        url = config.get(Config.SECTION_REMOTE_URL)
        if not url:
            return

        parsed = urlparse(url)

        user = parsed.username or config.get(Config.SECTION_REMOTE_USER)

        self.path_info = self.path_cls.from_parts(
            scheme=self.scheme,
            host=parsed.hostname,
            user=user,
            port=parsed.port,
            path=parsed.path,
        )

    def hdfs(self, path_info):
        import pyarrow

        return get_connection(
            pyarrow.hdfs.connect,
            path_info.host,
            path_info.port,
            user=path_info.user,
        )

    def hadoop_fs(self, cmd, user=None):
        cmd = "hadoop fs -" + cmd
        if user:
            cmd = "HADOOP_USER_NAME={} ".format(user) + cmd

        # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
        # See https://github.com/iterative/dvc/issues/1197.
        close_fds = os.name != "nt"

        executable = os.getenv("SHELL") if os.name != "nt" else None
        p = Popen(
            cmd,
            shell=True,
            close_fds=close_fds,
            executable=executable,
            env=fix_env(os.environ),
            stdin=PIPE,
            stdout=PIPE,
            stderr=PIPE,
        )
        out, err = p.communicate()
        if p.returncode != 0:
            raise RemoteCmdError(self.scheme, cmd, p.returncode, err)
        return out.decode("utf-8")

    @staticmethod
    def _group(regex, s, gname):
        match = re.match(regex, s)
        assert match is not None
        return match.group(gname)

    def get_file_checksum(self, path_info):
        # NOTE: pyarrow doesn't support checksum, so we need to use hadoop
        regex = r".*\t.*\t(?P<checksum>.*)"
        stdout = self.hadoop_fs(
            "checksum {}".format(path_info.path), user=path_info.user
        )
        return self._group(regex, stdout, "checksum")

    def copy(self, from_info, to_info, **_kwargs):
        dname = posixpath.dirname(to_info.path)
        with self.hdfs(to_info) as hdfs:
            hdfs.mkdir(dname)
            # NOTE: this is how `hadoop fs -cp` works too: it copies through
            # your local machine.
            with hdfs.open(from_info.path, "rb") as from_fobj:
                with hdfs.open(to_info.path, "wb") as to_fobj:
                    to_fobj.upload(from_fobj)

    def remove(self, path_info):
        if path_info.scheme != "hdfs":
            raise NotImplementedError

        if self.exists(path_info):
            logger.debug("Removing {}".format(path_info.path))
            with self.hdfs(path_info) as hdfs:
                hdfs.rm(path_info.path)

    def exists(self, path_info):
        assert not isinstance(path_info, list)
        assert path_info.scheme == "hdfs"
        with self.hdfs(path_info) as hdfs:
            return hdfs.exists(path_info.path)

    def _upload(self, from_file, to_info, **_kwargs):
        with self.hdfs(to_info) as hdfs:
            hdfs.mkdir(posixpath.dirname(to_info.path))
            tmp_file = tmp_fname(to_info.path)
            with open(from_file, "rb") as fobj:
                hdfs.upload(tmp_file, fobj)
            hdfs.rename(tmp_file, to_info.path)

    def _download(self, from_info, to_file, **_kwargs):
        with self.hdfs(from_info) as hdfs:
            with open(to_file, "wb+") as fobj:
                hdfs.download(from_info.path, fobj)

    @contextmanager
    def open(self, path_info, mode="r", encoding=None):
        assert mode in {"r", "rt", "rb"}

        try:
            with self.hdfs(path_info) as hdfs, closing(
                hdfs.open(path_info.path, mode="rb")
            ) as fd:
                if mode == "rb":
                    yield fd
                else:
                    yield io.TextIOWrapper(fd, encoding=encoding)
        except IOError as e:
            # Empty .errno and not specific enough error class in pyarrow,
            # see https://issues.apache.org/jira/browse/ARROW-6248
            if "file does not exist" in str(e):
                raise FileNotFoundError(*e.args)
            raise

    def list_cache_paths(self):
        if not self.exists(self.path_info):
            return []

        files = []
        dirs = [self.path_info.path]

        with self.hdfs(self.path_info) as hdfs:
            while dirs:
                for entry in hdfs.ls(dirs.pop(), detail=True):
                    if entry["kind"] == "directory":
                        dirs.append(urlparse(entry["name"]).path)
                    elif entry["kind"] == "file":
                        files.append(urlparse(entry["name"]).path)

        return files
<|endoftext|>


<commit_before>"""Image standardization.

See Nyul et al. 2000: New variants of a method of MRI scale standardization.
"""

from __future__ import division

import numpy as np

import dwi.files
import dwi.util

DEFAULT_CONFIGURATION = dict(
        pc=(0., 99.8),
        landmarks=[i*10 for i in range(1, 10)], # Deciles
        scale=(1, 4095),
        )

def landmark_scores(img, pc1, pc2, landmarks, thresholding=True):
    """Get scores at histogram landmarks."""
    from scipy.stats import scoreatpercentile
    if thresholding:
        threshold = np.mean(img)
        img = img[img > threshold]
    p1 = scoreatpercentile(img, pc1)
    p2 = scoreatpercentile(img, pc2)
    scores = [scoreatpercentile(img, i) for i in landmarks]
    return p1, p2, scores

def map_onto_scale(p1, p2, s1, s2, v):
    """Map value v from original scale [p1, p2] onto standard scale [s1, s2]."""
    assert p1 <= p2, (p1, p2)
    assert s1 <= s2, (s1, s2)
    if p1 == p2:
        assert s1 == s2, (s1, s2)
        return s1
    f = (v-p1) / (p2-p1)
    r = f * (s2-s1) + s1
    return r

def transform(img, p1, p2, scores, s1, s2, mapped_scores):
    """Transform image onto standard scale."""
    scores = [p1] + list(scores) + [p2]
    mapped_scores = [s1] + list(mapped_scores) + [s2]
    r = np.zeros_like(img, dtype=np.int)
    for pos, v in np.ndenumerate(img):
        # Select slot where to map.
        slot = sum(v > s for s in scores)
        slot = np.clip(slot, 1, len(scores)-1)
        r[pos] = map_onto_scale(scores[slot-1], scores[slot],
                mapped_scores[slot-1], mapped_scores[slot], v)
    return r

def standardize(img, cfg):
    """Transform an image based on a configuration (file)."""
    if isinstance(cfg, str) or isinstance(cfg, unicode):
        cfg = read_standardization_configuration(cfg)
    d = cfg
    p1, p2, scores = landmark_scores(img, d['pc1'], d['pc2'], d['landmarks'])
    img = transform(img, p1, p2, scores, d['s1'], d['s2'], d['mapped_scores'])
    return img

def write_standardization_configuration(filename, pc1, pc2, landmarks, s1, s2,
        mapped_scores):
    """Write image standardization configuration file."""
    with open(filename, 'w') as f:
        f.write(toline([pc1, pc2]))
        f.write(toline(landmarks))
        f.write(toline([s1, s2]))
        f.write(toline(mapped_scores))

def read_standardization_configuration(filename):
    """Read image standardization configuration file."""
    lines = list(dwi.files.valid_lines(filename))[:4]
    lines = [l.split() for l in lines]
    d = collections.OrderedDict()
    d['pc1'], d['pc2'] = map(float, lines[0])
    d['landmarks'] = map(float, lines[1])
    d['s1'], d['s2'] = map(int, lines[2])
    d['mapped_scores'] = map(int, lines[3])
    if len(d['landmarks']) != len(d['mapped_scores']):
        raise Exception('Invalid standardization file: {}'.format(filename))
    return d
<commit_msg>fixed call scope<commit_after>"""Image standardization.

See Nyul et al. 2000: New variants of a method of MRI scale standardization.
"""

from __future__ import division

import numpy as np

import dwi.files
import dwi.util

DEFAULT_CONFIGURATION = dict(
        pc=(0., 99.8),
        landmarks=[i*10 for i in range(1, 10)], # Deciles
        scale=(1, 4095),
        )

def landmark_scores(img, pc1, pc2, landmarks, thresholding=True):
    """Get scores at histogram landmarks."""
    from scipy.stats import scoreatpercentile
    if thresholding:
        threshold = np.mean(img)
        img = img[img > threshold]
    p1 = scoreatpercentile(img, pc1)
    p2 = scoreatpercentile(img, pc2)
    scores = [scoreatpercentile(img, i) for i in landmarks]
    return p1, p2, scores

def map_onto_scale(p1, p2, s1, s2, v):
    """Map value v from original scale [p1, p2] onto standard scale [s1, s2]."""
    assert p1 <= p2, (p1, p2)
    assert s1 <= s2, (s1, s2)
    if p1 == p2:
        assert s1 == s2, (s1, s2)
        return s1
    f = (v-p1) / (p2-p1)
    r = f * (s2-s1) + s1
    return r

def transform(img, p1, p2, scores, s1, s2, mapped_scores):
    """Transform image onto standard scale."""
    scores = [p1] + list(scores) + [p2]
    mapped_scores = [s1] + list(mapped_scores) + [s2]
    r = np.zeros_like(img, dtype=np.int)
    for pos, v in np.ndenumerate(img):
        # Select slot where to map.
        slot = sum(v > s for s in scores)
        slot = np.clip(slot, 1, len(scores)-1)
        r[pos] = map_onto_scale(scores[slot-1], scores[slot],
                mapped_scores[slot-1], mapped_scores[slot], v)
    return r

def standardize(img, cfg):
    """Transform an image based on a configuration (file)."""
    if isinstance(cfg, str) or isinstance(cfg, unicode):
        cfg = read_standardization_configuration(cfg)
    d = cfg
    p1, p2, scores = landmark_scores(img, d['pc1'], d['pc2'], d['landmarks'])
    img = transform(img, p1, p2, scores, d['s1'], d['s2'], d['mapped_scores'])
    return img

def write_standardization_configuration(filename, pc1, pc2, landmarks, s1, s2,
        mapped_scores):
    """Write image standardization configuration file."""
    with open(filename, 'w') as f:
        f.write(dwi.files.toline([pc1, pc2]))
        f.write(dwi.files.toline(landmarks))
        f.write(dwi.files.toline([s1, s2]))
        f.write(dwi.files.toline(mapped_scores))

def read_standardization_configuration(filename):
    """Read image standardization configuration file."""
    lines = list(dwi.files.valid_lines(filename))[:4]
    lines = [l.split() for l in lines]
    d = collections.OrderedDict()
    d['pc1'], d['pc2'] = map(float, lines[0])
    d['landmarks'] = map(float, lines[1])
    d['s1'], d['s2'] = map(int, lines[2])
    d['mapped_scores'] = map(int, lines[3])
    if len(d['landmarks']) != len(d['mapped_scores']):
        raise Exception('Invalid standardization file: {}'.format(filename))
    return d
<|endoftext|>


<commit_before>from optparse import make_option
import os
import os.path
import errno
import shutil

import bencode
from django.core.management.base import BaseCommand

from WhatManager2 import manage_torrent
from WhatManager2.utils import wm_unicode, wm_str
from home.models import WhatTorrent, DownloadLocation, ReplicaSet
from what_transcode.utils import get_info_hash

def safe_makedirs(p):
    try:
        os.makedirs(wm_str(p))
    except OSError as ex:
        if ex.errno != errno.EEXIST:
            raise ex

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
        make_option('--base-dir',
                    action='store_true',
                    dest='base_dir',
                    default=False,
                    help='Pass the containing directory of the torrent instead of the directory'
                         ' itself. The name of the torrent will be automatically appended.'),
    )
    help = u'Moves existing torrent data and import the torrent in WM.'

    def __init__(self):
        super(Command, self).__init__()
        self.pseudo_request = lambda: None
        self.trans_instance = ReplicaSet.get_what_master().get_preferred_instance()
        self.download_location = DownloadLocation.get_what_preferred()
        self.data_path = None
        self.torrent_path = None
        self.torrent_info = None
        self.info_hash = None
        self.dest_path = None
        self.what_torrent = None

    def check_args(self, args):
        if len(args) != 2:
            return False
        if not os.path.isdir(args[0]):
            return False
        if not os.path.isfile(args[1]):
            return False
        try:
            self.info_hash = get_info_hash(args[1])
        except Exception:
            print u'Invalid .torrent file.'
            return False
        return True

    def get_unicode_torrent_files(self):
        files = []
        for f in self.torrent_info['info']['files']:
            nf = dict(f)
            nf['path'] = [wm_unicode(i) for i in nf['path']]
            files.append(nf)
        return files

    def check_files(self):
        print u'Checking for existing files...'
        if 'files' in self.torrent_info['info']:
            for f in self.get_unicode_torrent_files():
                f_path = os.path.join(self.data_path, *f['path'])
                print wm_str(u'Checking {0}'.format(f_path))
                if not os.path.isfile(wm_str(f_path)):
                    print wm_str(u'{0} does not exist. What are you giving me?'.format(f_path))
                    return False
        else:
            f_path = os.path.join(self.data_path, self.torrent_info['info']['name'])
            print wm_str(u'Checking {0}'.format(f_path))
            if not os.path.isfile(wm_str(f_path)):
                print wm_str(u'{0} does not exist. What are you giving me?')
                return False
        print u'Creating destination directory...'
        self.dest_path = os.path.join(self.download_location.path, unicode(self.what_torrent.id))
        os.makedirs(wm_str(self.dest_path))
        os.chmod(self.dest_path, 0777)
        if 'files' in self.torrent_info['info']:
            self.dest_path = os.path.join(self.dest_path, wm_unicode(
                self.torrent_info['info']['name']))
            os.makedirs(wm_str(self.dest_path))
        print u'All torrent data files exist.'
        return True

    def move_files(self):
        print u'Moving files to new directory...'
        if 'files' in self.torrent_info['info']:
            for f in self.get_unicode_torrent_files():
                f_path = os.path.join(self.data_path, *f['path'])
                f_dest_path = os.path.join(self.dest_path, *f['path'])
                safe_makedirs(os.path.dirname(f_dest_path))
                shutil.move(wm_str(f_path), wm_str(f_dest_path))
        else:
            f_path = os.path.join(self.data_path, wm_unicode(self.torrent_info['info']['name']))
            f_dest_path = os.path.join(self.dest_path, wm_unicode(
                self.torrent_info['info']['name']))
            safe_makedirs(os.path.dirname(f_dest_path))
            shutil.move(wm_str(f_path), wm_str(f_dest_path))

    def handle(self, *args, **options):
        if not self.check_args(args):
            print u'Pass the torrent data directory as a first argument, ' \
                  u'a path to the .torrent file as a second.'
            return
        self.data_path, self.torrent_path = [wm_unicode(i) for i in args]
        with open(wm_str(self.torrent_path), 'rb') as f:
            self.torrent_info = bencode.bdecode(f.read())
        if options['base_dir']:
            self.data_path = os.path.join(self.data_path,
                                          wm_unicode(self.torrent_info['info']['name']))
        self.what_torrent = WhatTorrent.get_or_create(self.pseudo_request, info_hash=self.info_hash)
        if not self.check_files():
            return
        self.move_files()
        print 'Adding torrent to WM...'
        manage_torrent.add_torrent(self.pseudo_request, self.trans_instance,
                                   self.download_location, self.what_torrent.id)
        print 'Done!'
<commit_msg>Manual torrent import check<commit_after>from optparse import make_option
import os
import os.path
import errno
import shutil

import bencode
from django.core.management.base import BaseCommand

from WhatManager2 import manage_torrent
from WhatManager2.utils import wm_unicode, wm_str
from home.models import WhatTorrent, DownloadLocation, ReplicaSet, TransTorrent
from what_transcode.utils import get_info_hash

def safe_makedirs(p):
    try:
        os.makedirs(wm_str(p))
    except OSError as ex:
        if ex.errno != errno.EEXIST:
            raise ex

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
        make_option('--base-dir',
                    action='store_true',
                    dest='base_dir',
                    default=False,
                    help='Pass the containing directory of the torrent instead of the directory'
                         ' itself. The name of the torrent will be automatically appended.'),
    )
    help = u'Moves existing torrent data and import the torrent in WM.'

    def __init__(self):
        super(Command, self).__init__()
        self.pseudo_request = lambda: None
        self.trans_instance = ReplicaSet.get_what_master().get_preferred_instance()
        self.download_location = DownloadLocation.get_what_preferred()
        self.data_path = None
        self.torrent_path = None
        self.torrent_info = None
        self.info_hash = None
        self.dest_path = None
        self.what_torrent = None

    def check_args(self, args):
        if len(args) != 2:
            return False
        if not os.path.isdir(args[0]):
            return False
        if not os.path.isfile(args[1]):
            return False
        try:
            self.info_hash = get_info_hash(args[1])
        except Exception:
            print u'Invalid .torrent file.'
            return False
        return True

    def get_unicode_torrent_files(self):
        files = []
        for f in self.torrent_info['info']['files']:
            nf = dict(f)
            nf['path'] = [wm_unicode(i) for i in nf['path']]
            files.append(nf)
        return files

    def check_files(self):
        print u'Checking for existing files...'
        if 'files' in self.torrent_info['info']:
            for f in self.get_unicode_torrent_files():
                f_path = os.path.join(self.data_path, *f['path'])
                print wm_str(u'Checking {0}'.format(f_path))
                if not os.path.isfile(wm_str(f_path)):
                    print wm_str(u'{0} does not exist. What are you giving me?'.format(f_path))
                    return False
        else:
            f_path = os.path.join(self.data_path, self.torrent_info['info']['name'])
            print wm_str(u'Checking {0}'.format(f_path))
            if not os.path.isfile(wm_str(f_path)):
                print wm_str(u'{0} does not exist. What are you giving me?'.format(f_path))
                return False
        print u'Creating destination directory...'
        self.dest_path = os.path.join(self.download_location.path, unicode(self.what_torrent.id))
        os.makedirs(wm_str(self.dest_path))
        os.chmod(self.dest_path, 0777)
        if 'files' in self.torrent_info['info']:
            self.dest_path = os.path.join(self.dest_path, wm_unicode(
                self.torrent_info['info']['name']))
            os.makedirs(wm_str(self.dest_path))
        print u'All torrent data files exist.'
        return True

    def move_files(self):
        print u'Moving files to new directory...'
        if 'files' in self.torrent_info['info']:
            for f in self.get_unicode_torrent_files():
                f_path = os.path.join(self.data_path, *f['path'])
                f_dest_path = os.path.join(self.dest_path, *f['path'])
                safe_makedirs(os.path.dirname(f_dest_path))
                shutil.move(wm_str(f_path), wm_str(f_dest_path))
        else:
            f_path = os.path.join(self.data_path, wm_unicode(self.torrent_info['info']['name']))
            f_dest_path = os.path.join(self.dest_path, wm_unicode(
                self.torrent_info['info']['name']))
            safe_makedirs(os.path.dirname(f_dest_path))
            shutil.move(wm_str(f_path), wm_str(f_dest_path))

    def handle(self, *args, **options):
        if not self.check_args(args):
            print u'Pass the torrent data directory as a first argument, ' \
                  u'a path to the .torrent file as a second.'
            return
        self.data_path, self.torrent_path = [wm_unicode(i) for i in args]
        with open(wm_str(self.torrent_path), 'rb') as f:
            self.torrent_info = bencode.bdecode(f.read())
        if options['base_dir']:
            self.data_path = os.path.join(self.data_path,
                                          wm_unicode(self.torrent_info['info']['name']))
        print u'Checking to see if torrent is already loaded into WM..'
        masters = list(ReplicaSet.get_what_master().transinstance_set.all())
        try:
            TransTorrent.objects.get(instance__in=masters, info_hash=self.info_hash)
            print u'Torrent already added to WM. Skipping...'
            return False
        except TransTorrent.DoesNotExist:
            pass
        self.what_torrent = WhatTorrent.get_or_create(self.pseudo_request, info_hash=self.info_hash)
        if not self.check_files():
            return
        self.move_files()
        print 'Adding torrent to WM...'
        manage_torrent.add_torrent(self.pseudo_request, self.trans_instance,
                                   self.download_location, self.what_torrent.id)
        print 'Done!'
<|endoftext|>


<commit_before>#!/usr/bin/python

import roslib; roslib.load_manifest('rgbd_benchmark_tools')
import rospy
import rosbag
import sensor_msgs.msg
import argparse
import sys
import os
import cv
import struct
import array
import tf
import binascii

from cv_bridge import CvBridge, CvBridgeError

if __name__ == '__main__':

    # parse command line
    parser = argparse.ArgumentParser(description='''
    This scripts reads a bag file containing kinect RGBD images,
    adds the corresponding PointCloud2 messages, and saves it again into a bag file.
    Optional arguments allow to select only a portion of the original bag file.
    ''')
    parser.add_argument('--start', help='skip the first N seconds of  input bag file',default=0.00)
    parser.add_argument('--duration', help='only process N seconds of  input bag file')
    parser.add_argument('--nth', help='only process every N-th frame of input bag file',default=1)
    parser.add_argument('--skip', help='skip N blocks in the beginning', default=1)
    parser.add_argument('--nopoints', help='do not add point clouds', action='store_true')
    parser.add_argument('inputbag', help='input bag file')
    parser.add_argument('outputbag', nargs='?',help='output bag file')
    args = parser.parse_args()

    if not args.outputbag:
        args.outputbag = os.path.splitext(args.inputbag)[0] + "-points.bag"

    print "Processing bag file:"
    print "  in:",args.inputbag
    print "  out:",args.outputbag
    print "  starting from: %s seconds"%(args.start)

    if args.duration:
        print "  duration: %s seconds"%(args.duration)

    print "  saving every %s-th frame"%(args.nth)
    args.skip = float(args.skip)
    print "  skipping %s blocks"%(args.skip)

    inbag = rosbag.Bag(args.inputbag,'r')
    outbag = rosbag.Bag(args.outputbag, 'w', compression=rosbag.bag.Compression.BZ2)

    depth_camera_info = None
    rgb_camera_info = None
    depth_image = None
    rgb_image_color = None
    cortex = None

    nan = float('nan')
    bridge = CvBridge()
    frame = 0
    transforms = dict()

    time_start = None
    for topic, msg, t in inbag.read_messages():
        if time_start==None:
            time_start=t
        if t - time_start < rospy.Duration.from_sec(float(args.start)):
            continue
        if args.duration and (t - time_start > rospy.Duration.from_sec(float(args.start) + float(args.duration))):
            break
        print "t=%f\r"%(t-time_start).to_sec(),
        if topic == "/cortex_marker_array":
            cortex = msg
            continue
        if topic == "/tf":
            for transform in msg.transforms:
                transforms[ (transform.header.frame_id,transform.child_frame_id) ] = transform
            continue
        if topic == "/imu":
            imu = msg
            continue
        if topic == "/camera/depth/camera_info":
            depth_camera_info = msg
            continue
        if topic == "/camera/rgb/camera_info":
            rgb_camera_info = msg
            continue
        if topic == "/camera/rgb/image_color" and rgb_camera_info:
            rgb_image_color = msg
            continue
        if topic == "/camera/depth/image" and depth_camera_info and rgb_image_color and rgb_camera_info:
            depth_image = msg
            # now process frame

            frame += 1
            if frame % float(args.nth) ==0:
                if args.skip > 0:
                    args.skip -= 1
                    continue
                # store messages
                msg = tf.msg.tfMessage()
                msg.transforms = list( transforms.itervalues() )
                outbag.write("/tf",msg,t)
                if cortex:
                    outbag.write("/cortex_marker_array",cortex,t)
                transforms = dict()
                outbag.write("/imu",imu,t)
                outbag.write("/camera/depth/camera_info",depth_camera_info,t)
                outbag.write("/camera/depth/image",depth_image,t)
                outbag.write("/camera/rgb/camera_info",rgb_camera_info,t)
                outbag.write("/camera/rgb/image_color",rgb_image_color,t)
                # consume the images
                imu = None
                depth_image = None
                rgb_image_color = None
#                print depth_camera_info.header.stamp.to_sec()
#                print depth_image.header.stamp.to_sec()
#                print rgb_camera_info.header.stamp.to_sec()
#                print rgb_image_color.header.stamp.to_sec()

                # generate monochrome image from color image
#                cv_rgb_image_mono = bridge.imgmsg_to_cv(rgb_image_color, "mono8")
#                rgb_image_mono = bridge.cv_to_imgmsg(cv_rgb_image_mono)
#                rgb_image_mono.header = rgb_image_color.header
#                outbag.write("/camera/rgb/image_mono",rgb_image_mono,t)

                # generate depth and colored point cloud
                cv_depth_image = bridge.imgmsg_to_cv(depth_image, "passthrough")
                cv_rgb_image_color = bridge.imgmsg_to_cv(rgb_image_color, "bgr8")
                #cv.NamedWindow("image_scaled")
                #cv.ShowImage("image_scaled", cv_rgb_image_color)
                #cv.WaitKey(10)
#
                if args.nopoints:
                    continue
                centerX = depth_camera_info.K[2]
                centerY = depth_camera_info.K[5]
                depthFocalLength = depth_camera_info.K[0]
                depth_points = sensor_msgs.msg.PointCloud2()
                depth_points.header = depth_image.header
                depth_points.width = depth_image.width
                depth_points.height  = depth_image.height
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "x",offset = 0,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "y",offset = 4,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "z",offset = 8,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.point_step = 16
                depth_points.row_step = depth_points.point_step * depth_points.width
                buffer = []
                buffer_rgb = []
                for v in range(depth_image.height):
                    for u in range(depth_image.width):
                        d = cv_depth_image[v,u]
                        ptx = (u - centerX) * d / depthFocalLength;
                        pty = (v - centerY) * d / depthFocalLength;
                        ptz = d;
                        buffer.append(struct.pack('ffff',ptx,pty,ptz,1.0))
                depth_points.data = "".join(buffer)
                outbag.write("/camera/depth/points", depth_points, t)

                centerX = depth_camera_info.K[2]
                centerY = depth_camera_info.K[5]
                depthFocalLength = depth_camera_info.K[0]
                rgb_points = sensor_msgs.msg.PointCloud2()
                rgb_points.header = rgb_image_color.header
                rgb_points.width = depth_image.width
                rgb_points.height  = depth_image.height
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "x",offset = 0,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "y",offset = 4,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "z",offset = 8,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "rgb",offset = 16,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.point_step = 32
                rgb_points.row_step = rgb_points.point_step * rgb_points.width
                buffer = []
                for v in range(depth_image.height):
                    for u in range(depth_image.width):
                        d = cv_depth_image[v,u]
                        rgb = cv_rgb_image_color[v,u]
                        ptx = (u - centerX) * d / depthFocalLength;
                        pty = (v - centerY) * d / depthFocalLength;
                        ptz = d;
                        buffer.append(struct.pack('ffffBBBBIII',
                            ptx,pty,ptz,1.0,
                            rgb[0],rgb[1],rgb[2],0,
                            0,0,0))
                rgb_points.data = "".join(buffer)
                outbag.write("/camera/rgb/points", rgb_points, t)
            continue
        # anything else: pass thru
        outbag.write(topic,msg,t)

                #print frame

    outbag.close()
    print
<commit_msg>check for missing frames<commit_after>#!/usr/bin/python

import roslib; roslib.load_manifest('rgbd_benchmark_tools')
import rospy
import rosbag
import sensor_msgs.msg
import argparse
import sys
import os
import cv
import struct
import array
import tf
import binascii

from cv_bridge import CvBridge, CvBridgeError

if __name__ == '__main__':

    # parse command line
    parser = argparse.ArgumentParser(description='''
    This scripts reads a bag file containing kinect RGBD images,
    adds the corresponding PointCloud2 messages, and saves it again into a bag file.
    Optional arguments allow to select only a portion of the original bag file.
    ''')
    parser.add_argument('--start', help='skip the first N seconds of  input bag file',default=0.00)
    parser.add_argument('--duration', help='only process N seconds of  input bag file')
    parser.add_argument('--nth', help='only process every N-th frame of input bag file',default=1)
    parser.add_argument('--skip', help='skip N blocks in the beginning', default=1)
    parser.add_argument('--nopoints', help='do not add point clouds', action='store_true')
    parser.add_argument('inputbag', help='input bag file')
    parser.add_argument('outputbag', nargs='?',help='output bag file')
    args = parser.parse_args()

    if not args.outputbag:
        args.outputbag = os.path.splitext(args.inputbag)[0] + "-points.bag"

    print "Processing bag file:"
    print "  in:",args.inputbag
    print "  out:",args.outputbag
    print "  starting from: %s seconds"%(args.start)

    if args.duration:
        print "  duration: %s seconds"%(args.duration)

    print "  saving every %s-th frame"%(args.nth)
    args.skip = float(args.skip)
    print "  skipping %s blocks"%(args.skip)

    inbag = rosbag.Bag(args.inputbag,'r')
    outbag = rosbag.Bag(args.outputbag, 'w', compression=rosbag.bag.Compression.BZ2)

    depth_camera_info = None
    rgb_camera_info = None
    depth_image = None
    rgb_image_color = None
    cortex = None

    nan = float('nan')
    bridge = CvBridge()
    frame = 0
    transforms = dict()

    time_start = None
    for topic, msg, t in inbag.read_messages():
        if time_start==None:
            time_start=t
        if t - time_start < rospy.Duration.from_sec(float(args.start)):
            continue
        if args.duration and (t - time_start > rospy.Duration.from_sec(float(args.start) + float(args.duration))):
            break
        print "t=%f\r"%(t-time_start).to_sec(),
        if topic == "/cortex_marker_array":
            cortex = msg
            continue
        if topic == "/tf":
            for transform in msg.transforms:
                transforms[ (transform.header.frame_id,transform.child_frame_id) ] = transform
            continue
        if topic == "/imu":
            imu = msg
            continue
        if topic == "/camera/depth/camera_info":
            depth_camera_info = msg
            continue
        if topic == "/camera/rgb/camera_info":
            rgb_camera_info = msg
            continue
        if topic == "/camera/rgb/image_color" and rgb_camera_info:
            rgb_image_color = msg
            continue
        if topic == "/camera/depth/image" and depth_camera_info and rgb_image_color and rgb_camera_info:
            depth_image = msg
            # now process frame

            if depth_image.header.stamp - rgb_image_color.header.stamp > rospy.Duration.from_sec(1/30.0):
                continue

            frame += 1
            if frame % float(args.nth) ==0:
                if args.skip > 0:
                    args.skip -= 1
                    continue
                # store messages
                msg = tf.msg.tfMessage()
                msg.transforms = list( transforms.itervalues() )
                outbag.write("/tf",msg,t)
                if cortex:
                    outbag.write("/cortex_marker_array",cortex,t)
                transforms = dict()
                outbag.write("/imu",imu,t)
                outbag.write("/camera/depth/camera_info",depth_camera_info,t)
                outbag.write("/camera/depth/image",depth_image,t)
                outbag.write("/camera/rgb/camera_info",rgb_camera_info,t)
                outbag.write("/camera/rgb/image_color",rgb_image_color,t)
                # consume the images
                imu = None
                depth_image = None
                rgb_image_color = None
#                print depth_camera_info.header.stamp.to_sec()
#                print depth_image.header.stamp.to_sec()
#                print rgb_camera_info.header.stamp.to_sec()
#                print rgb_image_color.header.stamp.to_sec()

                # generate monochrome image from color image
#                cv_rgb_image_mono = bridge.imgmsg_to_cv(rgb_image_color, "mono8")
#                rgb_image_mono = bridge.cv_to_imgmsg(cv_rgb_image_mono)
#                rgb_image_mono.header = rgb_image_color.header
#                outbag.write("/camera/rgb/image_mono",rgb_image_mono,t)

                # generate depth and colored point cloud
                cv_depth_image = bridge.imgmsg_to_cv(depth_image, "passthrough")
                cv_rgb_image_color = bridge.imgmsg_to_cv(rgb_image_color, "bgr8")
                #cv.NamedWindow("image_scaled")
                #cv.ShowImage("image_scaled", cv_rgb_image_color)
                #cv.WaitKey(10)
#
                if args.nopoints:
                    continue
                centerX = depth_camera_info.K[2]
                centerY = depth_camera_info.K[5]
                depthFocalLength = depth_camera_info.K[0]
                depth_points = sensor_msgs.msg.PointCloud2()
                depth_points.header = depth_image.header
                depth_points.width = depth_image.width
                depth_points.height  = depth_image.height
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "x",offset = 0,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "y",offset = 4,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.fields.append(sensor_msgs.msg.PointField(
                    name = "z",offset = 8,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                depth_points.point_step = 16
                depth_points.row_step = depth_points.point_step * depth_points.width
                buffer = []
                buffer_rgb = []
                for v in range(depth_image.height):
                    for u in range(depth_image.width):
                        d = cv_depth_image[v,u]
                        ptx = (u - centerX) * d / depthFocalLength;
                        pty = (v - centerY) * d / depthFocalLength;
                        ptz = d;
                        buffer.append(struct.pack('ffff',ptx,pty,ptz,1.0))
                depth_points.data = "".join(buffer)
                outbag.write("/camera/depth/points", depth_points, t)

                centerX = depth_camera_info.K[2]
                centerY = depth_camera_info.K[5]
                depthFocalLength = depth_camera_info.K[0]
                rgb_points = sensor_msgs.msg.PointCloud2()
                rgb_points.header = rgb_image_color.header
                rgb_points.width = depth_image.width
                rgb_points.height  = depth_image.height
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "x",offset = 0,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "y",offset = 4,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "z",offset = 8,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.fields.append(sensor_msgs.msg.PointField(
                    name = "rgb",offset = 16,datatype = sensor_msgs.msg.PointField.FLOAT32,count = 1 ))
                rgb_points.point_step = 32
                rgb_points.row_step = rgb_points.point_step * rgb_points.width
                buffer = []
                for v in range(depth_image.height):
                    for u in range(depth_image.width):
                        d = cv_depth_image[v,u]
                        rgb = cv_rgb_image_color[v,u]
                        ptx = (u - centerX) * d / depthFocalLength;
                        pty = (v - centerY) * d / depthFocalLength;
                        ptz = d;
                        buffer.append(struct.pack('ffffBBBBIII',
                            ptx,pty,ptz,1.0,
                            rgb[0],rgb[1],rgb[2],0,
                            0,0,0))
                rgb_points.data = "".join(buffer)
                outbag.write("/camera/rgb/points", rgb_points, t)
            continue
        # anything else: pass thru
        outbag.write(topic,msg,t)

                #print frame

    outbag.close()
    print
<|endoftext|>


<commit_before>"""Module for generating annotated cells."""

import os
import random

import numpy as np
import scipy.misc
import scipy.ndimage

from jicbioimage.illustrate import AnnotatedImage

def post_process_annotation(ann, dilated_region, celldata, rotation, crop=True, rotate=True,
                            enlarge=True, padding=True):

    if crop:
        yis, xis = dilated_region.index_arrays
        ymin, ymax = np.min(yis), np.max(yis)
        xmin, xmax = np.min(xis), np.max(xis)
        ann = ann[ymin:ymax,
                  xmin:xmax]
        celldata["dy_offset"] = ymin
        celldata["dx_offset"] = xmin

    if padding:
        ydim, xdim, zdim = ann.shape
        p = 25
        pydim = ydim + p + p
        pxdim = xdim + p + p
        padded = AnnotatedImage.blank_canvas(width=pxdim, height=pydim)
        padded[p:ydim+p, p:xdim+p] = ann
        ann = padded
        celldata["dy_offset"] -= p
        celldata["dx_offset"] -= p
        celldata["ydim"] = pydim
        celldata["xdim"] = pxdim

    if enlarge:
        ann = scipy.misc.imresize(ann, 3.0, "nearest").view(AnnotatedImage)

    ann = scipy.ndimage.rotate(ann, rotation, order=0).view(AnnotatedImage)

    return ann

def write_cell_views(fpath_prefix, wall_projection, marker_projection, region, celldata,
                     crop=True, rotate=True, enlarge=True, padding=True):
    wall_ann = AnnotatedImage.from_grayscale(wall_projection, (1, 0, 0))
    marker_ann = AnnotatedImage.from_grayscale(marker_projection, (0, 1, 0))
    ann = wall_ann + marker_ann
    ann.mask_region(region.border, (200, 200, 200))
    dilated_region = region.dilate(20)
    wall_ann[np.logical_not(dilated_region)] = (0, 0, 0)
    marker_ann[np.logical_not(dilated_region)] = (0, 0, 0)
    ann[np.logical_not(dilated_region)] = (0, 0, 0)

    rotation = random.randrange(0, 360)
    celldata["rotation"] = rotation
    for suffix, annotation in [("-wall", wall_ann),
                               ("-marker", marker_ann),
                               ("-combined", ann)]:
        fpath = fpath_prefix + suffix + ".png"
        annotation = post_process_annotation(annotation, dilated_region, celldata,
                                             rotation, crop, enlarge, padding)

        scipy.misc.imsave(fpath, annotation)
<commit_msg>Made the rotation variant on 90 degrees<commit_after>"""Module for generating annotated cells."""

import os
import random

import numpy as np
import scipy.misc
import scipy.ndimage

from jicbioimage.illustrate import AnnotatedImage

def post_process_annotation(ann, dilated_region, celldata, rotation, crop=True, rotate=True,
                            enlarge=True, padding=True):

    if crop:
        yis, xis = dilated_region.index_arrays
        ymin, ymax = np.min(yis), np.max(yis)
        xmin, xmax = np.min(xis), np.max(xis)
        ann = ann[ymin:ymax,
                  xmin:xmax]
        celldata["dy_offset"] = ymin
        celldata["dx_offset"] = xmin

    if padding:
        ydim, xdim, zdim = ann.shape
        p = 25
        pydim = ydim + p + p
        pxdim = xdim + p + p
        padded = AnnotatedImage.blank_canvas(width=pxdim, height=pydim)
        padded[p:ydim+p, p:xdim+p] = ann
        ann = padded
        celldata["dy_offset"] -= p
        celldata["dx_offset"] -= p
        celldata["ydim"] = pydim
        celldata["xdim"] = pxdim

    if enlarge:
        ann = scipy.misc.imresize(ann, 3.0, "nearest").view(AnnotatedImage)

    ann = scipy.ndimage.rotate(ann, rotation, order=0).view(AnnotatedImage)

    return ann

def write_cell_views(fpath_prefix, wall_projection, marker_projection, region, celldata,
                     crop=True, rotate=True, enlarge=True, padding=True):
    wall_ann = AnnotatedImage.from_grayscale(wall_projection, (1, 0, 0))
    marker_ann = AnnotatedImage.from_grayscale(marker_projection, (0, 1, 0))
    ann = wall_ann + marker_ann
    ann.mask_region(region.border, (200, 200, 200))
    dilated_region = region.dilate(20)
    wall_ann[np.logical_not(dilated_region)] = (0, 0, 0)
    marker_ann[np.logical_not(dilated_region)] = (0, 0, 0)
    ann[np.logical_not(dilated_region)] = (0, 0, 0)

# If rotation is not 0, 90, 180, 270 the image becomes larger than then input
# and the scaling gets messed up. The scaling may not matter since downstream
# processing will use unit vectors around the center of the cell (the centroid)
# but for now I want to create annotations where the user clicks are reflected
# in the ouput.
#   rotation = random.randrange(0, 360)
    rotation = random.choice([0, 90, 180, 270])

    celldata["rotation"] = rotation
    for suffix, annotation in [("-wall", wall_ann),
                               ("-marker", marker_ann),
                               ("-combined", ann)]:
        fpath = fpath_prefix + suffix + ".png"
        annotation = post_process_annotation(annotation, dilated_region, celldata,
                                             rotation, crop, enlarge, padding)

        scipy.misc.imsave(fpath, annotation)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
#  Copyright (c) 2015, 2016, 2017 Tim Savannah under following terms:
#   You may modify and redistribe this script with your project
#
# It will download the latest GoodTests.py and use it to execute the tests.
#
#  This should be placed in a directory, "tests", at the root of your project. It assumes that ../$MY_PACKAGE_MODULE is the path to your test module, and will create a symlink to it in order to run tests.
#  The tests should be found in $MY_TEST_DIRECTORY in given "tests" folder.

# NOTE: Since version 1.2.3, you can also import this (like from a graphical application) and call the "main()" function.
#  All of the following globals are the defaults,  but can be overridden when calling main() (params have the same name as the globals).

import imp
import os

import glob
import subprocess
import sys

# URL to current version of GoodTests.py - You only need to change this if you host an internal copy.
GOODTESTS_URL = 'https://raw.githubusercontent.com/kata198/GoodTests/master/GoodTests.py'

# This should be your module name, and can be any relative or absolute path, or just a module name.
# If just a module name is given, the directory must be in current directory or parent directory.
MY_PACKAGE_MODULE = 'IndexedRedis'

#  Normally, you want to test the codebase during development, so you don't care about the site-packages installed version.
#     If you want to allow testing with any module by @MY_PACKAGE_MODULE in the python path, change this to True.
ALLOW_SITE_INSTALL = False

# This is the test directory that should contain all your tests. This should be a directory in your "tests" folder
MY_TEST_DIRECTORY = 'UnitTests'

__version__ = '2.1.0'
__version_tuple__ = (2, 1, 0)

def findGoodTests():
    '''
        findGoodTests - Tries to find GoodTests.py

        @return <dict> {
            'path' <str> -> Path to GoodTests.py (for execution)
            'success' <bool> -> True/False if we successfully found GoodTests.py
        }
    '''
    pathSplit = os.environ['PATH'].split(':')
    if '.' not in pathSplit:
        pathSplit = ['.'] + pathSplit
        os.environ['PATH'] = ':'.join(pathSplit)

    result = ''
    success = False
    for path in pathSplit:
        if path.endswith('/'):
            path = path[:-1]
        guess = path + '/GoodTests.py'
        if os.path.exists(guess):
            success = True
            result = guess
            break

    return {
        'path'    : result,
        "success" : success
    }

def download_goodTests(GOODTESTS_URL=None):
    '''
        download_goodTests - Attempts to download GoodTests, using the default global url (or one provided).

        @return <int> - 0 on success (program should continue), otherwise non-zero (program should abort with this exit status)
    '''
    if GOODTESTS_URL is None:
        GOODTESTS_URL = globals()['GOODTESTS_URL']

    validAnswer = False
    while validAnswer == False:
        sys.stdout.write('GoodTests not found. Would you like to install it to local folder? (y/n): ')
        sys.stdout.flush()
        answer = sys.stdin.readline().strip().lower()
        if answer not in ('y', 'n', 'yes', 'no'):
            continue
        validAnswer = True
        answer = answer[0]

    if answer == 'n':
        sys.stderr.write('Cannot run tests without installing GoodTests. http://pypi.python.org/pypi/GoodTests or https://github.com/kata198/Goodtests\n')
        return 1
    try:
        import urllib2 as urllib
    except ImportError:
        try:
            import urllib.request as urllib
        except:
            sys.stderr.write('Failed to import urllib. Trying pip.\n')
            import subprocess
            pipe = subprocess.Popen('pip install GoodTests', shell=True)
            res = pipe.wait()
            if res != 0:
                sys.stderr.write('Failed to install GoodTests with pip ordirect download. aborting.\n')
                return 1
    try:
        response = urllib.urlopen(GOODTESTS_URL)
        contents = response.read()
        if str != bytes:
            contents = contents.decode('ascii')
    except Exception as e:
        sys.stderr.write('Failed to download GoodTests.py from "%s"\n%s\n' %(GOODTESTS_URL, str(e)))
        return 1
    try:
        with open('GoodTests.py', 'w') as f:
            f.write(contents)
    except Exception as e:
        sys.stderr.write('Failed to write to GoodTests.py\n%s\n' %(str(e,)))
        return 1
    try:
        os.chmod('GoodTests.py', 0o775)
    except:
        sys.stderr.write('WARNING: Failed to chmod +x GoodTests.py, may not be able to be executed.\n')

    try:
        import GoodTests
    except ImportError:
        sys.stderr.write('Seemed to download GoodTests okay, but still cannot import. Aborting.\n')
        return 1

    return 0

def main(thisDir=None, additionalArgs=[], MY_PACKAGE_MODULE=None, ALLOW_SITE_INSTALL=None, MY_TEST_DIRECTORY=None, GOODTESTS_URL=None):
    '''
        Do the work - Try to find GoodTests.py, else prompt to download it, then run the tests.

        @param thisDir <None/str> - None to use default (directory this test file is in, or if not obtainable, current directory).
        @param additionalArgs <list> - Any additional args to pass to GoodTests.py

        Remainder of params take their global (top of file) defaults unless explicitly set here. See top of file for documentation.

        @return <int> - Exit code of application. 0 on success, non-zero on failure.

            TODO: Standardize return codes so external applications can derive failure without parsing error strings.
    '''

    if MY_PACKAGE_MODULE is None:
        MY_PACKAGE_MODULE = globals()['MY_PACKAGE_MODULE']
    if ALLOW_SITE_INSTALL is None:
        ALLOW_SITE_INSTALL = globals()['ALLOW_SITE_INSTALL']
    if MY_TEST_DIRECTORY is None:
        MY_TEST_DIRECTORY = globals()['MY_TEST_DIRECTORY']
    if GOODTESTS_URL is None:
        GOODTESTS_URL = globals()['GOODTESTS_URL']

    if not thisDir:
        thisDir = os.path.dirname(__file__)

    if not thisDir:
        thisDir = str(os.getcwd())
    elif not thisDir.startswith('/'):
        thisDir = str(os.getcwd()) + '/' + thisDir

    # If GoodTests is in current directory, make sure we find it later
    if os.path.exists('./GoodTests.py'):
        os.environ['PATH'] = str(os.getcwd()) + ':' + os.environ['PATH']

    os.chdir(thisDir)

    didDownload = False

    goodTestsInfo = findGoodTests()
    if goodTestsInfo['success'] is False:
        downloadRet = download_goodTests(GOODTESTS_URL)
        if downloadRet != 0:
            return downloadRet
        goodTestsInfo = findGoodTests()
        if goodTestsInfo['success'] is False:
            sys.stderr.write('Could not download or find GoodTests.py. Try to download it yourself using "pip install GoodTests", or wget %s\n' %( GOODTESTS_URL,))
            return 1

    PYTHON_ENDINGS = ('.py', '.pyc', '.pyo')

    origMyPackageModule = MY_PACKAGE_MODULE[:]
    baseName = os.path.basename(MY_PACKAGE_MODULE)
    dirName = os.path.dirname(MY_PACKAGE_MODULE)

    newPath = None
    if dirName not in ('.', ''):
        if dirName.startswith('.'):
            dirName = os.getcwd() + os.sep + dirName + os.sep
        newPath = dirName
    elif dirName == '':
        inCurrentDir = False
        try:
            imp.find_module(MY_PACKAGE_MODULE)
            inCurrentDir = True
        except ImportError:
            # COMPAT WITH PREVIOUS runTests.py: Try plain module in parent directory
            foundIt = False
            oldSysPath = sys.path[:]
            sys.path = [os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)]
            try:
                imp.find_module(MY_PACKAGE_MODULE)
                foundIt = True
                sys.path = oldSysPath
            except ImportError as e:
                sys.path = oldSysPath
                if not ALLOW_SITE_INSTALL:
                    sys.stderr.write('Cannot find "%s" locally.\n' %(MY_PACKAGE_MODULE,))
                    return 2
                else:
                    try:
                        __import__(baseName)
                    except:
                        sys.stderr.write('Cannot find "%s" locally or in global python path.\n' %(MY_PACKAGE_MODULE,))
                        return 2

            if foundIt is True:
                newPath = os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)
        if inCurrentDir is True:
            newPath = os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)

    if newPath:
        newPythonPath = [newPath] + [x for x in os.environ.get('PYTHONPATH', '').split(':') if x]
        os.environ['PYTHONPATH'] = ':'.join(newPythonPath)
        sys.path = [newPath] + sys.path

    try:
        __import__(baseName)
    except ImportError as e:
        if baseName.endswith(('.py', '.pyc', '.pyo')):
            MY_PACKAGE_MODULE = baseName[baseName.rindex('.')]

        if e.name != MY_PACKAGE_MODULE:
            sys.stderr.write('Error while importing %s: %s\n Likely this is another dependency that needs to be installed\nPerhaps run "pip install %s" or install the providing package.\n\n' %(e.name, str(e), e.name))
            return 1
        sys.stderr.write('Could not import %s. Either install it or otherwise add to PYTHONPATH\n%s\n' %(MY_PACKAGE_MODULE, str(e)))
        return 1

    if not os.path.isdir(MY_TEST_DIRECTORY):
        if not os.path.exists(MY_TEST_DIRECTORY):
            sys.stderr.write('Cannot find test directory: %s\n' %(MY_TEST_DIRECTORY,))
        else:
            sys.stderr.write('Provided test directory, "%s" is not a directory.\n' %(MY_TEST_DIRECTORY,))
        return 3

    sys.stdout.write('Starting test..\n')
    sys.stdout.flush()
    sys.stderr.flush()

    didTerminate = False

    pipe = subprocess.Popen([goodTestsInfo['path']] + additionalArgs + [MY_TEST_DIRECTORY], env=os.environ, shell=False)
    while True:
        try:
            pipe.wait()
            break
        except KeyboardInterrupt:
            if not didTerminate:
                pipe.terminate()
                didTerminate = True
            else:
                pipe.kill()
                break

    return 0

if __name__ == '__main__':
    ret = main(None, sys.argv[1:])
    sys.exit(ret)
<commit_msg>Update from ancient runTests.py 2.1.0 to new 3.0.5. Not going to be a release for this.<commit_after>#!/usr/bin/env python
#
#  Copyright (c) 2015, 2016, 2017 Tim Savannah under following terms:
#   You may modify and redistribe this script with your project
#
# It will download the latest GoodTests.py and use it to execute the tests.
#
#  This should be placed in a directory, "tests", at the root of your project. It assumes that ../$MY_PACKAGE_MODULE is the path to your test module, and will create a symlink to it in order to run tests.
#  The tests should be found in $MY_TEST_DIRECTORY in given "tests" folder.

# NOTE: Since version 1.2.3, you can also import this (like from a graphical application) and call the "main()" function.
#  All of the following globals are the defaults,  but can be overridden when calling main() (params have the same name as the globals).

# Assign a local function, "find_mod" to the interface to search
#  PYTHONPATH for importable module
try:
    # imp.find_module has been deprecated as of python 3.7, so
    #   prefer some alternate/newer interfaces first.
    import importlib

    try:
        # If we have the newest and therefore least-deprecated
        #   way, use it.
        _findModSpec = importlib.util.find_spec
        def find_mod(modName):
            '''
                find_mod - Find a module by name.

                  Similar to import #modName but only finds importable module,
                   does not actually import.

                 @raises ImportError on failure
            '''
            modSpec = _findModSpec(modName)
            if not modSpec:
                # imp.find_module raises import error if cannot find,
                #   but find_spec just returns None
                # So simulate the ImportError for common interface
                raise ImportError('No module named %s' %(modName, ))

            return modSpec

    except AttributeError:
        # We have importlib, but don't have importlib.util.find_spec

        # We could use importlib.import_module which is present in
        #  python 2.7, but that changes behaviour by actually
        #  importing (and thus additionally checking syntax/other).
        #
        # So just fall back to the old imp.find_module in this case

        try:
            # Clean up namespace
            del importlib
        except:
            pass
        # Fall back to imp.find_module implementation below
        raise ImportError('importlib but no importlib.util')
        #find_mod = lambda modName : importlib.import_module(modName)

except:
    # importlib is not present or has an unknown/dated interface,
    #   so fallback to the deprecated but oldest form
    import imp

    # Use a lambda to ensure only one arg is passed as that is
    #   our standard interface
    find_mod = lambda modName : imp.find_module(modName)

import os

import subprocess
import sys

# URL to current version of GoodTests.py - You only need to change this if you host an internal copy.
GOODTESTS_URL = 'https://raw.githubusercontent.com/kata198/GoodTests/master/GoodTests.py'

# This should be your module name, and can be any relative or absolute path, or just a module name.
# If just a module name is given, the directory must be in current directory or parent directory.
MY_PACKAGE_MODULE = 'IndexedRedis'

#  Normally, you want to test the codebase during development, so you don't care about the site-packages installed version.
#     If you want to allow testing with any module by @MY_PACKAGE_MODULE in the python path, change this to True.
ALLOW_SITE_INSTALL = False

# This is the test directory that should contain all your tests. This should be a directory in your "tests" folder
MY_TEST_DIRECTORY = 'UnitTests'

__version__ = '3.0.5'
__version_tuple__ = (3, 0, 5)

def findGoodTests():
    '''
        findGoodTests - Tries to find GoodTests.py

        @return <dict> {
            'path' <str> -> Path to GoodTests.py (for execution)
            'success' <bool> -> True/False if we successfully found GoodTests.py
        }
    '''
    pathSplit = os.environ['PATH'].split(':')
    if '.' not in pathSplit:
        pathSplit = ['.'] + pathSplit
        os.environ['PATH'] = ':'.join(pathSplit)

    result = ''
    success = False
    for path in pathSplit:
        if path.endswith('/'):
            path = path[:-1]
        guess = path + '/GoodTests.py'
        if os.path.exists(guess):
            success = True
            result = guess
            break

    return {
        'path'    : result,
        "success" : success
    }

def findExecutable(execName):
    '''
        findExecutable - Search PATH for an executable

        @return <dict> {
            'path' <str> -> Path to executable (if found, see "success")
            'success' <bool> -> True/False if we successfully found requested executable
        }
    '''

    pathSplit = os.environ['PATH'].split(':')
    if '.' not in pathSplit:
        pathSplit = ['.'] + pathSplit
        os.environ['PATH'] = ':'.join(pathSplit)

    result = ''
    success = False
    for path in pathSplit:
        if path.endswith(os.sep):
            path = path[:-1]
        guess = path + os.sep + execName
        if os.path.exists(guess):
            success = True
            result = guess
            break

    return {
        "path"    : result,
        "success" : success
    }

def findGoodTests():
    return findExecutable('GoodTests.py')

def try_pip_install():
    '''
        try to pip install GoodTests.py

        First, try via pip module.

        If that fails, try to locate pip by dirname(current python executable) + os.sep + pip
        If that does not exist, scan PATH for pip

        If found a valid pip executable, invoke it to install GoodTests
        otherwise, fail.
    '''

    didImport = False
    try:
        import pip
        didImport = True
    except:
        pass

    if didImport is True:
        print ( "Found pip as module=pip")
        res = pip.main(['install', 'GoodTests'])
        if res == 0:
            return 0
        sys.stderr.write('Failed to install GoodTests via pip module. Falling back to pip executable...\n\n')

    pipPath = os.path.dirname(sys.executable) + os.sep + 'pip'
    print ( 'Searching for pip at "%s"' %(pipPath, ) )
    if not os.path.exists(pipPath):
        print ( '"%s" does not exist. Scanning PATH to locate a usable pip executable' %(pipPath, ))
        pipPath = None
        searchResults = findExecutable('pip')
        if not searchResults['success']:
            sys.stderr.write('Failed to find a usable pip executable in PATH.\n')
            return 1 # Failed to locate a usable pip

        pipPath = searchResults['path']

    print ( 'Found pip executable at "%s"' %(pipPath, ) )
    print ( "Executing:  %s %s 'install' 'GoodTests'" %(sys.executable, pipPath) )
    pipe = subprocess.Popen([sys.executable, pipPath, 'install', 'GoodTests'], shell=False, env=os.environ)
    res = pipe.wait()

    return res

def download_goodTests(GOODTESTS_URL=None):
    '''
        download_goodTests - Attempts to download GoodTests, using the default global url (or one provided).

        @return <int> - 0 on success (program should continue), otherwise non-zero (program should abort with this exit status)
    '''
    if GOODTESTS_URL is None:
        GOODTESTS_URL = globals()['GOODTESTS_URL']

    validAnswer = False
    while validAnswer == False:
        sys.stdout.write('GoodTests not found. Would you like to install it to local folder? (y/n): ')
        sys.stdout.flush()
        answer = sys.stdin.readline().strip().lower()
        if answer not in ('y', 'n', 'yes', 'no'):
            continue
        validAnswer = True
        answer = answer[0]

    if answer == 'n':
        sys.stderr.write('Cannot run tests without installing GoodTests. http://pypi.python.org/pypi/GoodTests or https://github.com/kata198/Goodtests\n')
        return 1
    try:
        import urllib2 as urllib
    except ImportError:
        try:
            import urllib.request as urllib
        except:
            sys.stderr.write('Failed to import urllib. Trying pip.\n')
            res = try_pip_install()
            if res != 0:
                sys.stderr.write('Failed to install GoodTests with pip or direct download. aborting.\n')
                return 1
    try:
        response = urllib.urlopen(GOODTESTS_URL)
        contents = response.read()
        if str != bytes:
            contents = contents.decode('ascii')
    except Exception as e:
        sys.stderr.write('Failed to download GoodTests.py from "%s"\n%s\n' %(GOODTESTS_URL, str(e)))
        sys.stderr.write('\nTrying pip.\n')
        res = try_pip_install()
        if res != 0:
            sys.stderr.write('Failed to install GoodTests with pip or direct download. aborting.\n')
            return 1
    try:
        with open('GoodTests.py', 'w') as f:
            f.write(contents)
    except Exception as e:
        sys.stderr.write('Failed to write to GoodTests.py\n%s\n' %(str(e,)))
        return 1
    try:
        os.chmod('GoodTests.py', 0o775)
    except:
        sys.stderr.write('WARNING: Failed to chmod +x GoodTests.py, may not be able to be executed.\n')

    try:
        import GoodTests
    except ImportError:
        sys.stderr.write('Seemed to download GoodTests okay, but still cannot import. Aborting.\n')
        return 1

    return 0

def main(thisDir=None, additionalArgs=[], MY_PACKAGE_MODULE=None, ALLOW_SITE_INSTALL=None, MY_TEST_DIRECTORY=None, GOODTESTS_URL=None):
    '''
        Do the work - Try to find GoodTests.py, else prompt to download it, then run the tests.

        @param thisDir <None/str> - None to use default (directory this test file is in, or if not obtainable, current directory).
        @param additionalArgs <list> - Any additional args to pass to GoodTests.py

        Remainder of params take their global (top of file) defaults unless explicitly set here. See top of file for documentation.

        @return <int> - Exit code of application. 0 on success, non-zero on failure.

            TODO: Standardize return codes so external applications can derive failure without parsing error strings.
    '''

    if MY_PACKAGE_MODULE is None:
        MY_PACKAGE_MODULE = globals()['MY_PACKAGE_MODULE']
    if ALLOW_SITE_INSTALL is None:
        ALLOW_SITE_INSTALL = globals()['ALLOW_SITE_INSTALL']
    if MY_TEST_DIRECTORY is None:
        MY_TEST_DIRECTORY = globals()['MY_TEST_DIRECTORY']
    if GOODTESTS_URL is None:
        GOODTESTS_URL = globals()['GOODTESTS_URL']

    if not thisDir:
        thisDir = os.path.dirname(__file__)

    if not thisDir:
        thisDir = str(os.getcwd())
    elif not thisDir.startswith('/'):
        thisDir = str(os.getcwd()) + '/' + thisDir

    # If GoodTests is in current directory, make sure we find it later
    if os.path.exists('./GoodTests.py'):
        os.environ['PATH'] = str(os.getcwd()) + ':' + os.environ['PATH']

    os.chdir(thisDir)

    goodTestsInfo = findGoodTests()
    if goodTestsInfo['success'] is False:
        downloadRet = download_goodTests(GOODTESTS_URL)
        if downloadRet != 0:
            return downloadRet
        goodTestsInfo = findGoodTests()
        if goodTestsInfo['success'] is False:
            sys.stderr.write('Could not download or find GoodTests.py. Try to download it yourself using "pip install GoodTests", or wget %s\n' %( GOODTESTS_URL,))
            return 1

    baseName = os.path.basename(MY_PACKAGE_MODULE)
    dirName = os.path.dirname(MY_PACKAGE_MODULE)

    newPath = None
    if dirName not in ('.', ''):
        if dirName.startswith('.'):
            dirName = os.getcwd() + os.sep + dirName + os.sep
        newPath = dirName
    elif dirName == '':
        inCurrentDir = False
        try:
            find_mod(MY_PACKAGE_MODULE)
            inCurrentDir = True
        except ImportError:
            # COMPAT WITH PREVIOUS runTests.py: Try plain module in parent directory
            foundIt = False
            oldSysPath = sys.path[:]
            sys.path = [os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)]
            try:
                find_mod(MY_PACKAGE_MODULE)
                foundIt = True
                sys.path = oldSysPath
            except ImportError as e:
                sys.path = oldSysPath
                if not ALLOW_SITE_INSTALL:
                    sys.stderr.write('Cannot find "%s" locally.\n' %(MY_PACKAGE_MODULE,))
                    return 2
                else:
                    try:
                        __import__(baseName)
                    except:
                        sys.stderr.write('Cannot find "%s" locally or in global python path.\n' %(MY_PACKAGE_MODULE,))
                        return 2

            if foundIt is True:
                newPath = os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)
        if inCurrentDir is True:
            newPath = os.path.realpath(os.getcwd() + os.sep + '..' + os.sep)

    if newPath:
        newPythonPath = [newPath] + [x for x in os.environ.get('PYTHONPATH', '').split(':') if x]
        os.environ['PYTHONPATH'] = ':'.join(newPythonPath)
        sys.path = [newPath] + sys.path

    try:
        __import__(baseName)
    except ImportError as e:
        if baseName.endswith(('.py', '.pyc', '.pyo')):
            MY_PACKAGE_MODULE = baseName[ : baseName.rindex('.')]

        try:
            eName = e.name
        except AttributeError as noNameE:
            # Some platforms python2 does not have this attribute
            #   so pull it from the message
            eName = e.message.split()[-1]

        if eName != MY_PACKAGE_MODULE:
            sys.stderr.write('Error while importing %s: %s\n Likely this is another dependency that needs to be installed\nPerhaps run "pip install %s" or install the providing package.\n\n' %(eName, str(e), eName))
            return 1
        sys.stderr.write('Could not import %s. Either install it or otherwise add to PYTHONPATH\n%s\n' %(MY_PACKAGE_MODULE, str(e)))
        return 1

    if not os.path.isdir(MY_TEST_DIRECTORY):
        if not os.path.exists(MY_TEST_DIRECTORY):
            sys.stderr.write('Cannot find test directory: %s\n' %(MY_TEST_DIRECTORY,))
        else:
            sys.stderr.write('Provided test directory, "%s" is not a directory.\n' %(MY_TEST_DIRECTORY,))
        return 3

    sys.stdout.write('Starting test..\n')
    sys.stdout.flush()
    sys.stderr.flush()

    didTerminate = False
    pipe = subprocess.Popen([sys.executable, goodTestsInfo['path']] + additionalArgs + [MY_TEST_DIRECTORY], env=os.environ, shell=False)
    while True:
        try:
            pipe.wait()
            break
        except KeyboardInterrupt:
            if not didTerminate:
                pipe.terminate()
                didTerminate = True
            else:
                pipe.kill()
                break

    return 0

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