<commit_before>"""
Methods for exporting course data to XML
"""

import logging
import lxml.etree
from xblock.fields import Scope, Reference, ReferenceList, ReferenceValueDict
from xmodule.contentstore.content import StaticContent
from xmodule.exceptions import NotFoundError
from xmodule.modulestore import EdxJSONEncoder, ModuleStoreEnum
from xmodule.modulestore.inheritance import own_metadata
from fs.osfs import OSFS
from json import dumps
import json
import os
from path import path
import shutil
from xmodule.modulestore.draft_and_published import DIRECT_ONLY_CATEGORIES
from opaque_keys.edx.locator import CourseLocator

DRAFT_DIR = "drafts"
PUBLISHED_DIR = "published"
EXPORT_VERSION_FILE = "format.json"
EXPORT_VERSION_KEY = "export_format"

DEFAULT_CONTENT_FIELDS = ['metadata', 'data']

def export_to_xml(modulestore, contentstore, course_key, root_dir, course_dir):
    """
    Export all modules from `modulestore` and content from `contentstore` as xml to `root_dir`.

    `modulestore`: A `ModuleStore` object that is the source of the modules to export
    `contentstore`: A `ContentStore` object that is the source of the content to export, can be None
    `course_key`: The `CourseKey` of the `CourseModuleDescriptor` to export
    `root_dir`: The directory to write the exported xml to
    `course_dir`: The name of the directory inside `root_dir` to write the course content to
    """

    course = modulestore.get_course(course_key, depth=None)  # None means infinite
    fsm = OSFS(root_dir)
    export_fs = course.runtime.export_fs = fsm.makeopendir(course_dir)

    root = lxml.etree.Element('unknown')

    # export only the published content
    with modulestore.branch_setting(ModuleStoreEnum.Branch.published_only, course_key):
        # change all of the references inside the course to use the xml expected key type w/o version & branch
        xml_centric_course_key = CourseLocator(course_key.org, course_key.course, course_key.run, deprecated=True)
        adapt_references(course, xml_centric_course_key, export_fs)

        course.add_xml_to_node(root)

    with export_fs.open('course.xml', 'w') as course_xml:
        lxml.etree.ElementTree(root).write(course_xml)

    # export the static assets
    policies_dir = export_fs.makeopendir('policies')
    if contentstore:
        contentstore.export_all_for_course(
            course_key,
            root_dir + '/' + course_dir + '/static/',
            root_dir + '/' + course_dir + '/policies/assets.json',
        )

        # If we are using the default course image, export it to the
        # legacy location to support backwards compatibility.
        if course.course_image == course.fields['course_image'].default:
            try:
                course_image = contentstore.find(
                    StaticContent.compute_location(
                        course.id,
                        course.course_image
                    ),
                )
            except NotFoundError:
                pass
            else:
                output_dir = root_dir + '/' + course_dir + '/static/images/'
                if not os.path.isdir(output_dir):
                    os.makedirs(output_dir)
                with OSFS(output_dir).open('course_image.jpg', 'wb') as course_image_file:
                    course_image_file.write(course_image.data)

    # export the static tabs
    export_extra_content(export_fs, modulestore, xml_centric_course_key, 'static_tab', 'tabs', '.html')

    # export the custom tags
    export_extra_content(export_fs, modulestore, xml_centric_course_key, 'custom_tag_template', 'custom_tags')

    # export the course updates
    export_extra_content(export_fs, modulestore, xml_centric_course_key, 'course_info', 'info', '.html')

    # export the 'about' data (e.g. overview, etc.)
    export_extra_content(export_fs, modulestore, xml_centric_course_key, 'about', 'about', '.html')

    # export the grading policy
    course_run_policy_dir = policies_dir.makeopendir(course.location.name)
    with course_run_policy_dir.open('grading_policy.json', 'w') as grading_policy:
        grading_policy.write(dumps(course.grading_policy, cls=EdxJSONEncoder, sort_keys=True, indent=4))

    # export all of the course metadata in policy.json
    with course_run_policy_dir.open('policy.json', 'w') as course_policy:
        policy = {'course/' + course.location.name: own_metadata(course)}
        course_policy.write(dumps(policy, cls=EdxJSONEncoder, sort_keys=True, indent=4))

    #### DRAFTS ####
    # xml backed courses don't support drafts!
    if course.runtime.modulestore.get_modulestore_type() != ModuleStoreEnum.Type.xml:
        # NOTE: this code assumes that verticals are the top most draftable container
        # should we change the application, then this assumption will no longer be valid
        # NOTE: we need to explicitly implement the logic for setting the vertical's parent
        # and index here since the XML modulestore cannot load draft modules
        with modulestore.branch_setting(ModuleStoreEnum.Branch.draft_preferred, course_key):
            draft_verticals = modulestore.get_items(
                course_key,
                qualifiers={'category': 'vertical'},
                revision=ModuleStoreEnum.RevisionOption.draft_only
            )

            if len(draft_verticals) > 0:
                draft_course_dir = export_fs.makeopendir(DRAFT_DIR)
                for draft_vertical in draft_verticals:
                    parent_loc = modulestore.get_parent_location(
                        draft_vertical.location,
                        revision=ModuleStoreEnum.RevisionOption.draft_preferred
                    )
                    # Don't try to export orphaned items.
                    if parent_loc is not None:
                        logging.debug('parent_loc = {0}'.format(parent_loc))
                        if parent_loc.category in DIRECT_ONLY_CATEGORIES:
                            draft_vertical.xml_attributes['parent_sequential_url'] = parent_loc.to_deprecated_string()
                            sequential = modulestore.get_item(parent_loc)
                            index = sequential.children.index(draft_vertical.location)
                            draft_vertical.xml_attributes['index_in_children_list'] = str(index)
                        draft_vertical.runtime.export_fs = draft_course_dir
                        adapt_references(draft_vertical, xml_centric_course_key, draft_course_dir)
                        node = lxml.etree.Element('unknown')
                        draft_vertical.add_xml_to_node(node)

def adapt_references(subtree, destination_course_key, export_fs):
    """
    Map every reference in the subtree into destination_course_key and set it back into the xblock fields
    """
    subtree.runtime.export_fs = export_fs  # ensure everything knows where it's going!
    for field_name, field in subtree.fields.iteritems():
        if field.is_set_on(subtree):
            if isinstance(field, Reference):
                value = field.read_from(subtree)
                if value is not None:
                    field.write_to(subtree, field.read_from(subtree).map_into_course(destination_course_key))
            elif field_name == 'children':
                # don't change the children field but do recurse over the children
                [adapt_references(child, destination_course_key, export_fs) for child in subtree.get_children()]
            elif isinstance(field, ReferenceList):
                field.write_to(
                    subtree,
                    [ele.map_into_course(destination_course_key) for ele in field.read_from(subtree)]
                )
            elif isinstance(field, ReferenceValueDict):
                field.write_to(
                    subtree, {
                        key: ele.map_into_course(destination_course_key) for key, ele in field.read_from(subtree).iteritems()
                    }
                )

def _export_field_content(xblock_item, item_dir):
    """
    Export all fields related to 'xblock_item' other than 'metadata' and 'data' to json file in provided directory
    """
    module_data = xblock_item.get_explicitly_set_fields_by_scope(Scope.content)
    if isinstance(module_data, dict):
        for field_name in module_data:
            if field_name not in DEFAULT_CONTENT_FIELDS:
                # filename format: {dirname}.{field_name}.json
                with item_dir.open('{0}.{1}.{2}'.format(xblock_item.location.name, field_name, 'json'),
                                   'w') as field_content_file:
                    field_content_file.write(dumps(module_data.get(field_name, {}), cls=EdxJSONEncoder, sort_keys=True, indent=4))

def export_extra_content(export_fs, modulestore, course_key, category_type, dirname, file_suffix=''):
    items = modulestore.get_items(course_key, qualifiers={'category': category_type})

    if len(items) > 0:
        item_dir = export_fs.makeopendir(dirname)
        for item in items:
            adapt_references(item, course_key, export_fs)
            with item_dir.open(item.location.name + file_suffix, 'w') as item_file:
                item_file.write(item.data.encode('utf8'))

                # export content fields other then metadata and data in json format in current directory
                _export_field_content(item, item_dir)

def convert_between_versions(source_dir, target_dir):
    """
    Converts a version 0 export format to version 1, and vice versa.

    @param source_dir: the directory structure with the course export that should be converted.
       The contents of source_dir will not be altered.
    @param target_dir: the directory where the converted export should be written.
    @return: the version number of the converted export.
    """
    def convert_to_version_1():
        """ Convert a version 0 archive to version 0 """
        os.mkdir(copy_root)
        with open(copy_root / EXPORT_VERSION_FILE, 'w') as f:
            f.write('{{"{export_key}": 1}}\n'.format(export_key=EXPORT_VERSION_KEY))

        # If a drafts folder exists, copy it over.
        copy_drafts()

        # Now copy everything into the published directory
        published_dir = copy_root / PUBLISHED_DIR
        shutil.copytree(path(source_dir) / course_name, published_dir)
        # And delete the nested drafts directory, if it exists.
        nested_drafts_dir = published_dir / DRAFT_DIR
        if nested_drafts_dir.isdir():
            shutil.rmtree(nested_drafts_dir)

    def convert_to_version_0():
        """ Convert a version 1 archive to version 0 """
        # Copy everything in "published" up to the top level.
        published_dir = path(source_dir) / course_name / PUBLISHED_DIR
        if not published_dir.isdir():
            raise ValueError("a version 1 archive must contain a published branch")

        shutil.copytree(published_dir, copy_root)

        # If there is a DRAFT branch, copy it. All other branches are ignored.
        copy_drafts()

    def copy_drafts():
        """
        Copy drafts directory from the old archive structure to the new.
        """
        draft_dir = path(source_dir) / course_name / DRAFT_DIR
        if draft_dir.isdir():
            shutil.copytree(draft_dir, copy_root / DRAFT_DIR)

    root = os.listdir(source_dir)
    if len(root) != 1 or (path(source_dir) / root[0]).isfile():
        raise ValueError("source archive does not have single course directory at top level")

    course_name = root[0]

    # For this version of the script, we simply convert back and forth between version 0 and 1.
    original_version = get_version(path(source_dir) / course_name)
    if original_version not in [0, 1]:
        raise ValueError("unknown version: " + str(original_version))
    desired_version = 1 if original_version is 0 else 0

    copy_root = path(target_dir) / course_name

    if desired_version == 1:
        convert_to_version_1()
    else:
        convert_to_version_0()

    return desired_version

def get_version(course_path):
    """
    Return the export format version number for the given
    archive directory structure (represented as a path instance).

    If the archived file does not correspond to a known export
    format, None will be returned.
    """
    format_file = course_path / EXPORT_VERSION_FILE
    if not format_file.isfile():
        return 0
    with open(format_file, "r") as f:
        data = json.load(f)
        if EXPORT_VERSION_KEY in data:
            return data[EXPORT_VERSION_KEY]

    return None
<commit_msg>Turn on bulk operations for xml export<commit_after>"""
Methods for exporting course data to XML
"""

import logging
import lxml.etree
from xblock.fields import Scope, Reference, ReferenceList, ReferenceValueDict
from xmodule.contentstore.content import StaticContent
from xmodule.exceptions import NotFoundError
from xmodule.modulestore import EdxJSONEncoder, ModuleStoreEnum
from xmodule.modulestore.inheritance import own_metadata
from fs.osfs import OSFS
from json import dumps
import json
import os
from path import path
import shutil
from xmodule.modulestore.draft_and_published import DIRECT_ONLY_CATEGORIES
from opaque_keys.edx.locator import CourseLocator

DRAFT_DIR = "drafts"
PUBLISHED_DIR = "published"
EXPORT_VERSION_FILE = "format.json"
EXPORT_VERSION_KEY = "export_format"

DEFAULT_CONTENT_FIELDS = ['metadata', 'data']

def export_to_xml(modulestore, contentstore, course_key, root_dir, course_dir):
    """
    Export all modules from `modulestore` and content from `contentstore` as xml to `root_dir`.

    `modulestore`: A `ModuleStore` object that is the source of the modules to export
    `contentstore`: A `ContentStore` object that is the source of the content to export, can be None
    `course_key`: The `CourseKey` of the `CourseModuleDescriptor` to export
    `root_dir`: The directory to write the exported xml to
    `course_dir`: The name of the directory inside `root_dir` to write the course content to
    """

    with modulestore.bulk_operations(course_key):

        course = modulestore.get_course(course_key, depth=None)  # None means infinite
        fsm = OSFS(root_dir)
        export_fs = course.runtime.export_fs = fsm.makeopendir(course_dir)

        root = lxml.etree.Element('unknown')

        # export only the published content
        with modulestore.branch_setting(ModuleStoreEnum.Branch.published_only, course_key):
            # change all of the references inside the course to use the xml expected key type w/o version & branch
            xml_centric_course_key = CourseLocator(course_key.org, course_key.course, course_key.run, deprecated=True)
            adapt_references(course, xml_centric_course_key, export_fs)

            course.add_xml_to_node(root)

        with export_fs.open('course.xml', 'w') as course_xml:
            lxml.etree.ElementTree(root).write(course_xml)

        # export the static assets
        policies_dir = export_fs.makeopendir('policies')
        if contentstore:
            contentstore.export_all_for_course(
                course_key,
                root_dir + '/' + course_dir + '/static/',
                root_dir + '/' + course_dir + '/policies/assets.json',
            )

            # If we are using the default course image, export it to the
            # legacy location to support backwards compatibility.
            if course.course_image == course.fields['course_image'].default:
                try:
                    course_image = contentstore.find(
                        StaticContent.compute_location(
                            course.id,
                            course.course_image
                        ),
                    )
                except NotFoundError:
                    pass
                else:
                    output_dir = root_dir + '/' + course_dir + '/static/images/'
                    if not os.path.isdir(output_dir):
                        os.makedirs(output_dir)
                    with OSFS(output_dir).open('course_image.jpg', 'wb') as course_image_file:
                        course_image_file.write(course_image.data)

        # export the static tabs
        export_extra_content(export_fs, modulestore, xml_centric_course_key, 'static_tab', 'tabs', '.html')

        # export the custom tags
        export_extra_content(export_fs, modulestore, xml_centric_course_key, 'custom_tag_template', 'custom_tags')

        # export the course updates
        export_extra_content(export_fs, modulestore, xml_centric_course_key, 'course_info', 'info', '.html')

        # export the 'about' data (e.g. overview, etc.)
        export_extra_content(export_fs, modulestore, xml_centric_course_key, 'about', 'about', '.html')

        # export the grading policy
        course_run_policy_dir = policies_dir.makeopendir(course.location.name)
        with course_run_policy_dir.open('grading_policy.json', 'w') as grading_policy:
            grading_policy.write(dumps(course.grading_policy, cls=EdxJSONEncoder, sort_keys=True, indent=4))

        # export all of the course metadata in policy.json
        with course_run_policy_dir.open('policy.json', 'w') as course_policy:
            policy = {'course/' + course.location.name: own_metadata(course)}
            course_policy.write(dumps(policy, cls=EdxJSONEncoder, sort_keys=True, indent=4))

        #### DRAFTS ####
        # xml backed courses don't support drafts!
        if course.runtime.modulestore.get_modulestore_type() != ModuleStoreEnum.Type.xml:
            # NOTE: this code assumes that verticals are the top most draftable container
            # should we change the application, then this assumption will no longer be valid
            # NOTE: we need to explicitly implement the logic for setting the vertical's parent
            # and index here since the XML modulestore cannot load draft modules
            with modulestore.branch_setting(ModuleStoreEnum.Branch.draft_preferred, course_key):
                draft_verticals = modulestore.get_items(
                    course_key,
                    qualifiers={'category': 'vertical'},
                    revision=ModuleStoreEnum.RevisionOption.draft_only
                )

                if len(draft_verticals) > 0:
                    draft_course_dir = export_fs.makeopendir(DRAFT_DIR)
                    for draft_vertical in draft_verticals:
                        parent_loc = modulestore.get_parent_location(
                            draft_vertical.location,
                            revision=ModuleStoreEnum.RevisionOption.draft_preferred
                        )
                        # Don't try to export orphaned items.
                        if parent_loc is not None:
                            logging.debug('parent_loc = {0}'.format(parent_loc))
                            if parent_loc.category in DIRECT_ONLY_CATEGORIES:
                                draft_vertical.xml_attributes['parent_sequential_url'] = parent_loc.to_deprecated_string()
                                sequential = modulestore.get_item(parent_loc)
                                index = sequential.children.index(draft_vertical.location)
                                draft_vertical.xml_attributes['index_in_children_list'] = str(index)
                            draft_vertical.runtime.export_fs = draft_course_dir
                            adapt_references(draft_vertical, xml_centric_course_key, draft_course_dir)
                            node = lxml.etree.Element('unknown')
                            draft_vertical.add_xml_to_node(node)

def adapt_references(subtree, destination_course_key, export_fs):
    """
    Map every reference in the subtree into destination_course_key and set it back into the xblock fields
    """
    subtree.runtime.export_fs = export_fs  # ensure everything knows where it's going!
    for field_name, field in subtree.fields.iteritems():
        if field.is_set_on(subtree):
            if isinstance(field, Reference):
                value = field.read_from(subtree)
                if value is not None:
                    field.write_to(subtree, field.read_from(subtree).map_into_course(destination_course_key))
            elif field_name == 'children':
                # don't change the children field but do recurse over the children
                [adapt_references(child, destination_course_key, export_fs) for child in subtree.get_children()]
            elif isinstance(field, ReferenceList):
                field.write_to(
                    subtree,
                    [ele.map_into_course(destination_course_key) for ele in field.read_from(subtree)]
                )
            elif isinstance(field, ReferenceValueDict):
                field.write_to(
                    subtree, {
                        key: ele.map_into_course(destination_course_key) for key, ele in field.read_from(subtree).iteritems()
                    }
                )

def _export_field_content(xblock_item, item_dir):
    """
    Export all fields related to 'xblock_item' other than 'metadata' and 'data' to json file in provided directory
    """
    module_data = xblock_item.get_explicitly_set_fields_by_scope(Scope.content)
    if isinstance(module_data, dict):
        for field_name in module_data:
            if field_name not in DEFAULT_CONTENT_FIELDS:
                # filename format: {dirname}.{field_name}.json
                with item_dir.open('{0}.{1}.{2}'.format(xblock_item.location.name, field_name, 'json'),
                                   'w') as field_content_file:
                    field_content_file.write(dumps(module_data.get(field_name, {}), cls=EdxJSONEncoder, sort_keys=True, indent=4))

def export_extra_content(export_fs, modulestore, course_key, category_type, dirname, file_suffix=''):
    items = modulestore.get_items(course_key, qualifiers={'category': category_type})

    if len(items) > 0:
        item_dir = export_fs.makeopendir(dirname)
        for item in items:
            adapt_references(item, course_key, export_fs)
            with item_dir.open(item.location.name + file_suffix, 'w') as item_file:
                item_file.write(item.data.encode('utf8'))

                # export content fields other then metadata and data in json format in current directory
                _export_field_content(item, item_dir)

def convert_between_versions(source_dir, target_dir):
    """
    Converts a version 0 export format to version 1, and vice versa.

    @param source_dir: the directory structure with the course export that should be converted.
       The contents of source_dir will not be altered.
    @param target_dir: the directory where the converted export should be written.
    @return: the version number of the converted export.
    """
    def convert_to_version_1():
        """ Convert a version 0 archive to version 0 """
        os.mkdir(copy_root)
        with open(copy_root / EXPORT_VERSION_FILE, 'w') as f:
            f.write('{{"{export_key}": 1}}\n'.format(export_key=EXPORT_VERSION_KEY))

        # If a drafts folder exists, copy it over.
        copy_drafts()

        # Now copy everything into the published directory
        published_dir = copy_root / PUBLISHED_DIR
        shutil.copytree(path(source_dir) / course_name, published_dir)
        # And delete the nested drafts directory, if it exists.
        nested_drafts_dir = published_dir / DRAFT_DIR
        if nested_drafts_dir.isdir():
            shutil.rmtree(nested_drafts_dir)

    def convert_to_version_0():
        """ Convert a version 1 archive to version 0 """
        # Copy everything in "published" up to the top level.
        published_dir = path(source_dir) / course_name / PUBLISHED_DIR
        if not published_dir.isdir():
            raise ValueError("a version 1 archive must contain a published branch")

        shutil.copytree(published_dir, copy_root)

        # If there is a DRAFT branch, copy it. All other branches are ignored.
        copy_drafts()

    def copy_drafts():
        """
        Copy drafts directory from the old archive structure to the new.
        """
        draft_dir = path(source_dir) / course_name / DRAFT_DIR
        if draft_dir.isdir():
            shutil.copytree(draft_dir, copy_root / DRAFT_DIR)

    root = os.listdir(source_dir)
    if len(root) != 1 or (path(source_dir) / root[0]).isfile():
        raise ValueError("source archive does not have single course directory at top level")

    course_name = root[0]

    # For this version of the script, we simply convert back and forth between version 0 and 1.
    original_version = get_version(path(source_dir) / course_name)
    if original_version not in [0, 1]:
        raise ValueError("unknown version: " + str(original_version))
    desired_version = 1 if original_version is 0 else 0

    copy_root = path(target_dir) / course_name

    if desired_version == 1:
        convert_to_version_1()
    else:
        convert_to_version_0()

    return desired_version

def get_version(course_path):
    """
    Return the export format version number for the given
    archive directory structure (represented as a path instance).

    If the archived file does not correspond to a known export
    format, None will be returned.
    """
    format_file = course_path / EXPORT_VERSION_FILE
    if not format_file.isfile():
        return 0
    with open(format_file, "r") as f:
        data = json.load(f)
        if EXPORT_VERSION_KEY in data:
            return data[EXPORT_VERSION_KEY]

    return None
<|endoftext|>


<commit_before>import cmd
import getpass
import grp
import inspect
import os
import pwd
import re
import readline
import shlex
import stat
import sys
import textwrap
import urlparse
import zlib

from errors import WebHDFSError
from client import WebHDFSClient
from attrib import LocalFSObject

# Work around python's overeager completer delimiters
readline.set_completer_delims(readline.get_completer_delims().translate(None, '-'))

class WebHDFSPrompt(cmd.Cmd):
    def __init__(self, base, conf=None, path=None, task=None, wait=None):
        cmd.Cmd.__init__(self)

        self.base = urlparse.urlparse(base)
        self.user = getpass.getuser()
        self.hdfs = WebHDFSClient(base, self.user, conf, wait)

        self.do_cd(path)

        if task:
            self.onecmd(' '.join(task))
            sys.exit(0)

    def _list_dir(self, sources):
        subdirs = []
        objects = []
        columns = ['mode', 'repl', 'owner', 'group', 'size', 'date', 'name']
        lengths = dict(zip(columns, [0] * len(columns)))
        build = {
            'date': '{:%b %d %Y %H:%M:%S}',
        }
        align = {
            'repl': '>',
            'size': '>',
        }

        for item in sources:
            if not stat.S_ISREG(item.perm) and not stat.S_ISDIR(item.perm):
                continue

            if stat.S_ISDIR(item.perm):
                subdirs.append(item.name)

            tmp_obj = {}

            for name in columns:
                text = build.get(name, '{}').format(getattr(item, name))

                tmp_obj[name] = text
                lengths[name] = max(lengths[name], len(text))

            objects.append(tmp_obj)

        text = ' '.join('{%s:%s%s}' % (i, align.get(i, ''), lengths[i]) for i in columns)
        for item in objects:
            print text.format(**item)

        return subdirs

    def _fix_path(self, path, local=False, required=False):
        path = '' if path is None else path.strip()
        rval = []

        if not path and required:
            raise WebHDFSError('%s: path not specified' % required)

        if not path:
            path = getattr(self, 'path', '/user/%s' % self.user) if not local else os.getcwd()

        if not path.startswith('/'):
            path = '%s/%s' % (self.path if not local else os.getcwd(), path)

        for part in path.split('/'):
            if not part or part == '.':
                continue
            if rval and part == '..':
                rval.pop()
            else:
                rval.append(part)

        return '/'+'/'.join(rval)

    def _print_usage(self):
        print getattr(self, inspect.stack()[1][3]).__doc__.strip().split('\n')[0]

    def _reset_prompt(self):
        self.prompt = '%s@%s r:%s l:%s> ' % (self.user, self.base.netloc, self.path, os.getcwd())

    def _complete_local(self, part, kind):
        path = self._fix_path(part, local=True)
        name = ''

        if part and not part.endswith('/'):
            name = os.path.basename(path)
            path = os.path.dirname(path)

        if kind == 'file':
            pick = lambda x: x.startswith(name) and not stat.S_ISDIR(os.stat('%s/%s' % (path, x)).st_mode)
        elif kind == 'dir':
            pick = lambda x: x.startswith(name) and stat.S_ISDIR(os.stat('%s/%s' % (path, x)).st_mode)
        else:
            pick = lambda x: x.startswith(name)

        return [i + ('/' if stat.S_ISDIR(os.stat('%s/%s' % (path, i)).st_mode) else ' ') for i in os.listdir(path) if pick(i)]

    def _complete_remote(self, part, kind):
        path = self._fix_path(part)
        name = ''

        if part and not part.endswith('/') and not part.endswith('/..'):
            name = os.path.basename(path)
            path = os.path.dirname(path)

        if kind == 'file':
            pick = lambda x: x.name.startswith(name) and not x.is_dir()
        elif kind == 'dir':
            pick = lambda x: x.name.startswith(name) and x.is_dir()
        else:
            pick = lambda x: x.name.startswith(name)

        return [i.name + ('/' if i.is_dir() else ' ') for i in self.hdfs.ls(path, request=pick)]

    def _complete_ownership(self, part, cache={}):
        if ':' not in part:
            if 'pwd' not in cache:
                cache['pwd'] = pwd.getpwall()

            return [i.pw_name for i in cache['pwd'] if i.pw_name.startswith(part)]
        else:
            if 'grp' not in cache:
                cache['grp'] = grp.getgrall()

            return [i.gr_name for i in cache['grp'] if i.gr_name.startswith(part.split(':', 1)[-1])]

    def completedefault(self, part, line, s, e):
        if part == '.' or part == '..':
            return [part + '/']

        args = shlex.split(line[:e])
        if len(args) == 1 or line[e - 1] == ' ':
            args.append('')

        # Extract completion magic from method documentation
        docs = getattr(getattr(self, 'do_'+args[0], object), '__doc__')
        rule = re.findall(r'(?:<(.+?)>)+', docs)[len(args) - 2]

        if re.search(r'(?:local|remote) (?:file/dir|file|dir)', rule):
            kind, dest = rule.split()
            return getattr(self, '_complete_'+kind)(args[-1], dest)
        if rule == '[owner][:group]':
            return self._complete_ownership(args[-1])

    def emptyline(self):
        pass

    def default(self, arg):
        print '%s: unknown command' % arg

    def do_cd(self, path=None):
        '''
            Usage: cd <remote dir>

            Changes the shell remote directory
        '''
        try:
            path = self._fix_path(path or '/user/%s' % self.user)
            if not self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: not a directory' % path)
            self.path = path
        except WebHDFSError as e:
            self.path = '/'
            print e
        finally:
            self._reset_prompt()

    def do_lcd(self, path=None):
        '''
            Usage: lcd <local dir>

            Changes the shell local directory
        '''
        try:
            path = self._fix_path(path or pwd.getpwnam(self.user).pw_dir, local=True)
            os.chdir(path)
        except (KeyError, OSError) as e:
            print e
        finally:
            self._reset_prompt()

    def do_ls(self, path=None):
        '''
            Usage: ls <remote file/dir>

            Lists remote file or directory
        '''
        try:
            path = self._fix_path(path)
            self._list_dir(self.hdfs.ls(path))
        except WebHDFSError as e:
            print e

    def do_lsr(self, path=None):
        '''
            Usage: ls <remote file/dir>

            Lists remote file or directory recursively
        '''
        try:
            path = self._fix_path(path)
            print path + ':'
            for name in self._list_dir(self.hdfs.ls(path)):
                print
                self.do_lsr('%s/%s' % (path, name))
        except WebHDFSError as e:
            print e

    def do_glob(self, path=None):
        '''
            Usage: glob <remote file/dir>

            Lists remote file or directory pattern
        '''
        try:
            path = self._fix_path(path, required='glob')
            self._list_dir(self.hdfs.glob(path))
        except WebHDFSError as e:
            print e

    def do_lls(self, path=None):
        '''
            Usage: lls <local file/dir>

            Lists local file or directory
        '''
        try:
            path = self._fix_path(path, local=True)
            info = os.stat(path)
            objs = []

            if stat.S_ISDIR(info.st_mode):
                objs = list(LocalFSObject(path, name) for name in os.listdir(path))
            elif stat.S_ISREG(info.st_mode):
                objs = [LocalFSObject(os.path.dirname(path), os.path.basename(path))]

            self._list_dir(objs)
        except OSError as e:
            print e

    def do_du(self, path=None):
        '''
            Usage: du <remote file/dir>

            Displays disk usage for remote file or directory
        '''
        try:
            path = self._fix_path(path)
            print self.hdfs.du(path)
        except WebHDFSError as e:
            print e

    def do_mkdir(self, path):
        '''
            Usage: mkdir <remote dir>

            Creates remote directory
        '''
        try:
            path = self._fix_path(path, required='mkdir')
            if self.hdfs.stat(path, catch=True):
                raise WebHDFSError('%s: already exists' % path)
            self.hdfs.mkdir(path)
        except WebHDFSError as e:
            print e

    def do_mv(self, args):
        '''
            Usage: mv <remote file/dir> <remote dir>

            Moves/renames remote file or directory
        '''
        try:
            path, dest = shlex.split(args)
            path = self._fix_path(path, required='mv')
            dest = self._fix_path(dest, required='mv')
            stat = self.hdfs.stat(dest, catch=True) or self.hdfs.stat(os.path.dirname(dest), catch=True)
            if stat and not stat.is_dir():
                raise WebHDFSError('%s: invalid destination' % dest)
            if not self.hdfs.mv(path, dest):
                raise WebHDFSError('%s: failed to move/rename' % path)
        except WebHDFSError as e:
            print e
        except ValueError as e:
            self._print_usage()

    def do_rm(self, path):
        '''
            Usage: rm <remote file>

            Removes remote file
        '''
        try:
            path = self._fix_path(path, required='rm')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot remove directory' % path)
            self.hdfs.rm(path)
        except WebHDFSError as e:
            print e

    def do_rmdir(self, path):
        '''
            Usage: rm <remote dir>

            Removes remote directory
        '''
        try:
            path = self._fix_path(path, required='rmdir')
            temp = self.hdfs.stat(path)
            if not temp.is_dir():
                raise WebHDFSError('%s: not a directory' % path)
            if not temp.is_empty():
                raise WebHDFSError('%s: directory not empty' % path)
            self.hdfs.rm(path)
        except WebHDFSError as e:
            print e

    def do_chown(self, args):
        '''
            Usage: chown <[owner][:group]> <remote file/dir>

            Change ownership of remote file to owner and/or group
        '''
        try:
            dest, path = shlex.split(args)
            path = self._fix_path(path, required='chown')
            o, g = dest.split(':', 1) if ':' in dest else (dest, '')
            self.hdfs.chown(path, owner=o, group=g)
        except WebHDFSError as e:
            print e
        except ValueError:
            self._print_usage()

    def do_get(self, path):
        '''
            Usage: get <remote file>

            Fetch remote file into current local directory
        '''
        try:
            path = self._fix_path(path, required='get')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot download directory' % path)
            if os.path.exists(os.path.basename(path)):
                raise WebHDFSError('%s: file exists' % path)
            self.hdfs.get(path, data=open('%s/%s' % (os.getcwd(), os.path.basename(path)), 'w'))
        except (WebHDFSError, OSError) as e:
            print e

    def do_put(self, path):
        '''
            Usage: put <local file>

            Upload local file into current remote directory
        '''
        try:
            path = self._fix_path(path, local=True, required='put')
            dest = '%s/%s' % (self.path, os.path.basename(path))
            if stat.S_ISDIR(os.stat(path).st_mode):
                raise WebHDFSError('%s: cannot upload directory' % path)
            if self.hdfs.stat(dest, catch=True):
                raise WebHDFSError('%s: already exists' % dest)
            self.hdfs.put(dest, data=open(path, 'r'))
        except (WebHDFSError, OSError) as e:
            print e

    def do_cat(self, path):
        '''
            Usage: cat <remote file>

            Display contents of remote file
        '''
        try:
            path = self._fix_path(path, required='cat')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot cat directory' % path)
            sys.stdout.write(self.hdfs.get(path))
        except (WebHDFSError, OSError) as e:
            print e

    def do_zcat(self, path):
        '''
            Usage: zcat <remote file>

            Display contents of compressed remote file
        '''
        try:
            path = self._fix_path(path, required='zcat')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot cat directory' % path)
            sys.stdout.write(zlib.decompress(self.hdfs.get(path), 16 + zlib.MAX_WBITS))
        except (WebHDFSError, OSError) as e:
            print e

    def do_EOF(self, line):
        print
        return True

for name, func in vars(WebHDFSPrompt).items():
    if name.startswith('do_') and getattr(func, '__doc__'):
        func.__doc__ = textwrap.dedent(func.__doc__).strip()
<commit_msg>Issue #36: Skip to the next parameter when there's only one completion candidate.<commit_after>import cmd
import getpass
import grp
import inspect
import os
import pwd
import re
import readline
import shlex
import stat
import sys
import textwrap
import urlparse
import zlib

from errors import WebHDFSError
from client import WebHDFSClient
from attrib import LocalFSObject

# Work around python's overeager completer delimiters
readline.set_completer_delims(readline.get_completer_delims().translate(None, '-'))

class WebHDFSPrompt(cmd.Cmd):
    def __init__(self, base, conf=None, path=None, task=None, wait=None):
        cmd.Cmd.__init__(self)

        self.base = urlparse.urlparse(base)
        self.user = getpass.getuser()
        self.hdfs = WebHDFSClient(base, self.user, conf, wait)

        self.do_cd(path)

        if task:
            self.onecmd(' '.join(task))
            sys.exit(0)

    def _list_dir(self, sources):
        subdirs = []
        objects = []
        columns = ['mode', 'repl', 'owner', 'group', 'size', 'date', 'name']
        lengths = dict(zip(columns, [0] * len(columns)))
        build = {
            'date': '{:%b %d %Y %H:%M:%S}',
        }
        align = {
            'repl': '>',
            'size': '>',
        }

        for item in sources:
            if not stat.S_ISREG(item.perm) and not stat.S_ISDIR(item.perm):
                continue

            if stat.S_ISDIR(item.perm):
                subdirs.append(item.name)

            tmp_obj = {}

            for name in columns:
                text = build.get(name, '{}').format(getattr(item, name))

                tmp_obj[name] = text
                lengths[name] = max(lengths[name], len(text))

            objects.append(tmp_obj)

        text = ' '.join('{%s:%s%s}' % (i, align.get(i, ''), lengths[i]) for i in columns)
        for item in objects:
            print text.format(**item)

        return subdirs

    def _fix_path(self, path, local=False, required=False):
        path = '' if path is None else path.strip()
        rval = []

        if not path and required:
            raise WebHDFSError('%s: path not specified' % required)

        if not path:
            path = getattr(self, 'path', '/user/%s' % self.user) if not local else os.getcwd()

        if not path.startswith('/'):
            path = '%s/%s' % (self.path if not local else os.getcwd(), path)

        for part in path.split('/'):
            if not part or part == '.':
                continue
            if rval and part == '..':
                rval.pop()
            else:
                rval.append(part)

        return '/'+'/'.join(rval)

    def _print_usage(self):
        print getattr(self, inspect.stack()[1][3]).__doc__.strip().split('\n')[0]

    def _reset_prompt(self):
        self.prompt = '%s@%s r:%s l:%s> ' % (self.user, self.base.netloc, self.path, os.getcwd())

    def _complete_local(self, part, kind):
        path = self._fix_path(part, local=True)
        name = ''

        if part and not part.endswith('/'):
            name = os.path.basename(path)
            path = os.path.dirname(path)

        if kind == 'file':
            pick = lambda x: x.startswith(name) and not stat.S_ISDIR(os.stat('%s/%s' % (path, x)).st_mode)
        elif kind == 'dir':
            pick = lambda x: x.startswith(name) and stat.S_ISDIR(os.stat('%s/%s' % (path, x)).st_mode)
        else:
            pick = lambda x: x.startswith(name)

        return [i + ('/' if stat.S_ISDIR(os.stat('%s/%s' % (path, i)).st_mode) else ' ') for i in os.listdir(path) if pick(i)]

    def _complete_remote(self, part, kind):
        path = self._fix_path(part)
        name = ''

        if part and not part.endswith('/') and not part.endswith('/..'):
            name = os.path.basename(path)
            path = os.path.dirname(path)

        if kind == 'file':
            pick = lambda x: x.name.startswith(name) and not x.is_dir()
        elif kind == 'dir':
            pick = lambda x: x.name.startswith(name) and x.is_dir()
        else:
            pick = lambda x: x.name.startswith(name)

        return [i.name + ('/' if i.is_dir() and not i.is_empty() else ' ') for i in self.hdfs.ls(path, request=pick)]

    def _complete_ownership(self, part, cache={}):
        if ':' not in part:
            if 'pwd' not in cache:
                cache['pwd'] = pwd.getpwall()

            rval = [i.pw_name for i in cache['pwd'] if i.pw_name.startswith(part)]
            return rval if len(rval) != 1 else [rval[0] + ':']
        else:
            if 'grp' not in cache:
                cache['grp'] = grp.getgrall()

            rval = [i.gr_name for i in cache['grp'] if i.gr_name.startswith(part.split(':', 1)[-1])]
            return rval if len(rval) != 1 else [rval[0] + ' ']

    def completedefault(self, part, line, s, e):
        if part == '.' or part == '..':
            return [part + '/']

        args = shlex.split(line[:e])
        if len(args) == 1 or line[e - 1] == ' ':
            args.append('')

        # Extract completion magic from method documentation
        docs = getattr(getattr(self, 'do_'+args[0], object), '__doc__')
        rule = re.findall(r'(?:<(.+?)>)+', docs)[len(args) - 2]

        if re.search(r'(?:local|remote) (?:file/dir|file|dir)', rule):
            kind, dest = rule.split()
            return getattr(self, '_complete_'+kind)(args[-1], dest)
        if rule == '[owner][:group]':
            return self._complete_ownership(args[-1])

    def emptyline(self):
        pass

    def default(self, arg):
        print '%s: unknown command' % arg

    def do_cd(self, path=None):
        '''
            Usage: cd <remote dir>

            Changes the shell remote directory
        '''
        try:
            path = self._fix_path(path or '/user/%s' % self.user)
            if not self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: not a directory' % path)
            self.path = path
        except WebHDFSError as e:
            self.path = '/'
            print e
        finally:
            self._reset_prompt()

    def do_lcd(self, path=None):
        '''
            Usage: lcd <local dir>

            Changes the shell local directory
        '''
        try:
            path = self._fix_path(path or pwd.getpwnam(self.user).pw_dir, local=True)
            os.chdir(path)
        except (KeyError, OSError) as e:
            print e
        finally:
            self._reset_prompt()

    def do_ls(self, path=None):
        '''
            Usage: ls <remote file/dir>

            Lists remote file or directory
        '''
        try:
            path = self._fix_path(path)
            self._list_dir(self.hdfs.ls(path))
        except WebHDFSError as e:
            print e

    def do_lsr(self, path=None):
        '''
            Usage: ls <remote file/dir>

            Lists remote file or directory recursively
        '''
        try:
            path = self._fix_path(path)
            print path + ':'
            for name in self._list_dir(self.hdfs.ls(path)):
                print
                self.do_lsr('%s/%s' % (path, name))
        except WebHDFSError as e:
            print e

    def do_glob(self, path=None):
        '''
            Usage: glob <remote file/dir>

            Lists remote file or directory pattern
        '''
        try:
            path = self._fix_path(path, required='glob')
            self._list_dir(self.hdfs.glob(path))
        except WebHDFSError as e:
            print e

    def do_lls(self, path=None):
        '''
            Usage: lls <local file/dir>

            Lists local file or directory
        '''
        try:
            path = self._fix_path(path, local=True)
            info = os.stat(path)
            objs = []

            if stat.S_ISDIR(info.st_mode):
                objs = list(LocalFSObject(path, name) for name in os.listdir(path))
            elif stat.S_ISREG(info.st_mode):
                objs = [LocalFSObject(os.path.dirname(path), os.path.basename(path))]

            self._list_dir(objs)
        except OSError as e:
            print e

    def do_du(self, path=None):
        '''
            Usage: du <remote file/dir>

            Displays disk usage for remote file or directory
        '''
        try:
            path = self._fix_path(path)
            print self.hdfs.du(path)
        except WebHDFSError as e:
            print e

    def do_mkdir(self, path):
        '''
            Usage: mkdir <remote dir>

            Creates remote directory
        '''
        try:
            path = self._fix_path(path, required='mkdir')
            if self.hdfs.stat(path, catch=True):
                raise WebHDFSError('%s: already exists' % path)
            self.hdfs.mkdir(path)
        except WebHDFSError as e:
            print e

    def do_mv(self, args):
        '''
            Usage: mv <remote file/dir> <remote dir>

            Moves/renames remote file or directory
        '''
        try:
            path, dest = shlex.split(args)
            path = self._fix_path(path, required='mv')
            dest = self._fix_path(dest, required='mv')
            stat = self.hdfs.stat(dest, catch=True) or self.hdfs.stat(os.path.dirname(dest), catch=True)
            if stat and not stat.is_dir():
                raise WebHDFSError('%s: invalid destination' % dest)
            if not self.hdfs.mv(path, dest):
                raise WebHDFSError('%s: failed to move/rename' % path)
        except WebHDFSError as e:
            print e
        except ValueError as e:
            self._print_usage()

    def do_rm(self, path):
        '''
            Usage: rm <remote file>

            Removes remote file
        '''
        try:
            path = self._fix_path(path, required='rm')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot remove directory' % path)
            self.hdfs.rm(path)
        except WebHDFSError as e:
            print e

    def do_rmdir(self, path):
        '''
            Usage: rm <remote dir>

            Removes remote directory
        '''
        try:
            path = self._fix_path(path, required='rmdir')
            temp = self.hdfs.stat(path)
            if not temp.is_dir():
                raise WebHDFSError('%s: not a directory' % path)
            if not temp.is_empty():
                raise WebHDFSError('%s: directory not empty' % path)
            self.hdfs.rm(path)
        except WebHDFSError as e:
            print e

    def do_chown(self, args):
        '''
            Usage: chown <[owner][:group]> <remote file/dir>

            Change ownership of remote file or directory
        '''
        try:
            dest, path = shlex.split(args)
            path = self._fix_path(path, required='chown')
            o, g = dest.split(':', 1) if ':' in dest else (dest, '')
            self.hdfs.chown(path, owner=o, group=g)
        except WebHDFSError as e:
            print e
        except ValueError:
            self._print_usage()

    def do_get(self, path):
        '''
            Usage: get <remote file>

            Fetch remote file into current local directory
        '''
        try:
            path = self._fix_path(path, required='get')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot download directory' % path)
            if os.path.exists(os.path.basename(path)):
                raise WebHDFSError('%s: file exists' % path)
            self.hdfs.get(path, data=open('%s/%s' % (os.getcwd(), os.path.basename(path)), 'w'))
        except (WebHDFSError, OSError) as e:
            print e

    def do_put(self, path):
        '''
            Usage: put <local file>

            Upload local file into current remote directory
        '''
        try:
            path = self._fix_path(path, local=True, required='put')
            dest = '%s/%s' % (self.path, os.path.basename(path))
            if stat.S_ISDIR(os.stat(path).st_mode):
                raise WebHDFSError('%s: cannot upload directory' % path)
            if self.hdfs.stat(dest, catch=True):
                raise WebHDFSError('%s: already exists' % dest)
            self.hdfs.put(dest, data=open(path, 'r'))
        except (WebHDFSError, OSError) as e:
            print e

    def do_cat(self, path):
        '''
            Usage: cat <remote file>

            Display contents of remote file
        '''
        try:
            path = self._fix_path(path, required='cat')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot cat directory' % path)
            sys.stdout.write(self.hdfs.get(path))
        except (WebHDFSError, OSError) as e:
            print e

    def do_zcat(self, path):
        '''
            Usage: zcat <remote file>

            Display contents of compressed remote file
        '''
        try:
            path = self._fix_path(path, required='zcat')
            if self.hdfs.stat(path).is_dir():
                raise WebHDFSError('%s: cannot cat directory' % path)
            sys.stdout.write(zlib.decompress(self.hdfs.get(path), 16 + zlib.MAX_WBITS))
        except (WebHDFSError, OSError) as e:
            print e

    def do_EOF(self, line):
        print
        return True

for name, func in vars(WebHDFSPrompt).items():
    if name.startswith('do_') and getattr(func, '__doc__'):
        func.__doc__ = textwrap.dedent(func.__doc__).strip()
<|endoftext|>


<commit_before><commit_msg>[DF] Retrieve dataset name and files through Internal::TreeUtils functions<commit_after><|endoftext|>


<commit_before>#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from bigdl.util.common import Sample as BSample, JTensor as BJTensor,\
    JavaCreator, _get_gateway, _py2java, _java2py
import numpy as np
import os
import tempfile
import uuid

from urllib.parse import urlparse

def convert_to_safe_path(input_path, follow_symlinks=True):
    # resolves symbolic links
    if follow_symlinks:
        return os.path.realpath(input_path)
    # covert to abs path
    return os.path.abspath(input_path)

def to_list_of_numpy(elements):
    if isinstance(elements, np.ndarray):
        return [elements]
    elif np.isscalar(elements):
        return [np.array(elements)]
    elif not isinstance(elements, list):
        raise ValueError("Wrong type: %s" % type(elements))

    results = []
    for element in elements:
        if np.isscalar(element):
            results.append(np.array(element))
        elif isinstance(element, np.ndarray):
            results.append(element)
        else:
            raise ValueError("Wrong type: %s" % type(element))

    return results

def get_file_list(path, recursive=False):
    return callZooFunc("float", "listPaths", path, recursive)

def is_local_path(path):
    parse_result = urlparse(path)
    return len(parse_result.scheme.lower()) == 0 or parse_result.scheme.lower() == "file"

def append_suffix(prefix, path):
    # append suffix
    splits = path.split(".")
    if len(splits) > 0:
        file_name = prefix + "." + splits[-1]
    else:
        file_name = prefix

    return file_name

def save_file(save_func, path, **kwargs):

    if is_local_path(path):
        save_func(path, **kwargs)
    else:
        file_name = str(uuid.uuid1())
        file_name = append_suffix(file_name, path)
        temp_path = os.path.join(tempfile.gettempdir(), file_name)

        try:
            save_func(temp_path, **kwargs)
            if "overwrite" in kwargs:
                put_local_file_to_remote(temp_path, path, over_write=kwargs['overwrite'])
            else:
                put_local_file_to_remote(temp_path, path)
        finally:
            os.remove(temp_path)

def load_from_file(load_func, path):
    if is_local_path(path):
        return load_func(path)
    else:
        file_name = str(uuid.uuid1())
        file_name = append_suffix(file_name, path)
        temp_path = os.path.join(tempfile.gettempdir(), file_name)
        get_remote_file_to_local(path, temp_path)
        try:
            return load_func(temp_path)
        finally:
            os.remove(temp_path)

def get_remote_file_to_local(remote_path, local_path, over_write=False):
    callZooFunc("float", "getRemoteFileToLocal", remote_path, local_path, over_write)

def put_local_file_to_remote(local_path, remote_path, over_write=False):
    callZooFunc("float", "putLocalFileToRemote", local_path, remote_path, over_write)

def set_core_number(num):
    callZooFunc("float", "setCoreNumber", num)

def callZooFunc(bigdl_type, name, *args):
    """ Call API in PythonBigDL """
    gateway = _get_gateway()
    args = [_py2java(gateway, a) for a in args]
    error = Exception("Cannot find function: %s" % name)
    for jinvoker in JavaCreator.instance(bigdl_type, gateway).value:
        # hasattr(jinvoker, name) always return true here,
        # so you need to invoke the method to check if it exist or not
        try:
            api = getattr(jinvoker, name)
            java_result = api(*args)
            result = _java2py(gateway, java_result)
        except Exception as e:
            error = e
            if not ("does not exist" in str(e)
                    and "Method {}".format(name) in str(e)):
                raise e
        else:
            return result
    raise error

class JTensor(BJTensor):

    def __init__(self, storage, shape, bigdl_type="float", indices=None):
        super(JTensor, self).__init__(storage, shape, bigdl_type, indices)

    @classmethod
    def from_ndarray(cls, a_ndarray, bigdl_type="float"):
        """
        Convert a ndarray to a DenseTensor which would be used in Java side.
        """
        if a_ndarray is None:
            return None
        assert isinstance(a_ndarray, np.ndarray), \
            "input should be a np.ndarray, not %s" % type(a_ndarray)
        return cls(a_ndarray,
                   a_ndarray.shape,
                   bigdl_type)

class Sample(BSample):

    def __init__(self, features, labels, bigdl_type="float"):
        super(Sample, self).__init__(features, labels, bigdl_type)

    @classmethod
    def from_ndarray(cls, features, labels, bigdl_type="float"):
        features = to_list_of_numpy(features)
        labels = to_list_of_numpy(labels)
        return cls(
            features=[JTensor(feature, feature.shape) for feature in features],
            labels=[JTensor(label, label.shape) for label in labels],
            bigdl_type=bigdl_type)
<commit_msg>Support save & load from hdfs in orca estimators (#4283)<commit_after>#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from bigdl.util.common import Sample as BSample, JTensor as BJTensor,\
    JavaCreator, _get_gateway, _py2java, _java2py
import numpy as np
import os
import tempfile
import uuid
import functools

from urllib.parse import urlparse

def convert_to_safe_path(input_path, follow_symlinks=True):
    # resolves symbolic links
    if follow_symlinks:
        return os.path.realpath(input_path)
    # covert to abs path
    return os.path.abspath(input_path)

def to_list_of_numpy(elements):
    if isinstance(elements, np.ndarray):
        return [elements]
    elif np.isscalar(elements):
        return [np.array(elements)]
    elif not isinstance(elements, list):
        raise ValueError("Wrong type: %s" % type(elements))

    results = []
    for element in elements:
        if np.isscalar(element):
            results.append(np.array(element))
        elif isinstance(element, np.ndarray):
            results.append(element)
        else:
            raise ValueError("Wrong type: %s" % type(element))

    return results

def get_file_list(path, recursive=False):
    return callZooFunc("float", "listPaths", path, recursive)

def is_local_path(path):
    parse_result = urlparse(path)
    return len(parse_result.scheme.lower()) == 0 or parse_result.scheme.lower() == "file"

def append_suffix(prefix, path):
    # append suffix
    splits = path.split(".")
    if len(splits) > 0:
        file_name = prefix + "." + splits[-1]
    else:
        file_name = prefix

    return file_name

def enable_multi_fs_save(save_func):

    @functools.wraps(save_func)
    def save_mult_fs(obj, path, *args, **kwargs):
        if is_local_path(path):
            return save_func(obj, path, *args, **kwargs)
        else:
            file_name = str(uuid.uuid1())
            file_name = append_suffix(file_name, path)
            temp_path = os.path.join(tempfile.gettempdir(), file_name)

            try:
                result = save_func(obj, temp_path, *args, **kwargs)
                if "overwrite" in kwargs:
                    put_local_file_to_remote(temp_path, path, over_write=kwargs['overwrite'])
                else:
                    put_local_file_to_remote(temp_path, path)
            finally:
                os.remove(temp_path)
            return result

    return save_mult_fs

def enable_multi_fs_load_static(load_func):
    @functools.wraps(load_func)
    def multi_fs_load(path, *args, **kwargs):
        if is_local_path(path):
            return load_func(path, *args, **kwargs)
        else:
            file_name = str(uuid.uuid1())
            file_name = append_suffix(file_name, path)
            temp_path = os.path.join(tempfile.gettempdir(), file_name)
            get_remote_file_to_local(path, temp_path)
            try:
                return load_func(temp_path, *args, **kwargs)
            finally:
                os.remove(temp_path)

    return multi_fs_load

def enable_multi_fs_load(load_func):

    @functools.wraps(load_func)
    def multi_fs_load(obj, path, *args, **kwargs):
        if is_local_path(path):
            return load_func(obj, path, *args, **kwargs)
        else:
            file_name = str(uuid.uuid1())
            file_name = append_suffix(file_name, path)
            temp_path = os.path.join(tempfile.gettempdir(), file_name)
            get_remote_file_to_local(path, temp_path)
            try:
                return load_func(obj, temp_path, *args, **kwargs)
            finally:
                os.remove(temp_path)

    return multi_fs_load

def get_remote_file_to_local(remote_path, local_path, over_write=False):
    callZooFunc("float", "getRemoteFileToLocal", remote_path, local_path, over_write)

def put_local_file_to_remote(local_path, remote_path, over_write=False):
    callZooFunc("float", "putLocalFileToRemote", local_path, remote_path, over_write)

def set_core_number(num):
    callZooFunc("float", "setCoreNumber", num)

def callZooFunc(bigdl_type, name, *args):
    """ Call API in PythonBigDL """
    gateway = _get_gateway()
    args = [_py2java(gateway, a) for a in args]
    error = Exception("Cannot find function: %s" % name)
    for jinvoker in JavaCreator.instance(bigdl_type, gateway).value:
        # hasattr(jinvoker, name) always return true here,
        # so you need to invoke the method to check if it exist or not
        try:
            api = getattr(jinvoker, name)
            java_result = api(*args)
            result = _java2py(gateway, java_result)
        except Exception as e:
            error = e
            if not ("does not exist" in str(e)
                    and "Method {}".format(name) in str(e)):
                raise e
        else:
            return result
    raise error

class JTensor(BJTensor):

    def __init__(self, storage, shape, bigdl_type="float", indices=None):
        super(JTensor, self).__init__(storage, shape, bigdl_type, indices)

    @classmethod
    def from_ndarray(cls, a_ndarray, bigdl_type="float"):
        """
        Convert a ndarray to a DenseTensor which would be used in Java side.
        """
        if a_ndarray is None:
            return None
        assert isinstance(a_ndarray, np.ndarray), \
            "input should be a np.ndarray, not %s" % type(a_ndarray)
        return cls(a_ndarray,
                   a_ndarray.shape,
                   bigdl_type)

class Sample(BSample):

    def __init__(self, features, labels, bigdl_type="float"):
        super(Sample, self).__init__(features, labels, bigdl_type)

    @classmethod
    def from_ndarray(cls, features, labels, bigdl_type="float"):
        features = to_list_of_numpy(features)
        labels = to_list_of_numpy(labels)
        return cls(
            features=[JTensor(feature, feature.shape) for feature in features],
            labels=[JTensor(label, label.shape) for label in labels],
            bigdl_type=bigdl_type)
<|endoftext|>


<commit_before>import sys
sys.path.append('..')

import os
import json
from time import time
import numpy as np
from tqdm import tqdm
from matplotlib import pyplot as plt
from sklearn.externals import joblib

import theano
import theano.tensor as T
from theano.sandbox.cuda.dnn import dnn_conv

from lib import activations
from lib import updates
from lib import inits
from lib.vis import color_grid_vis
from lib.rng import py_rng, np_rng
from lib.ops import batchnorm, conv_cond_concat, deconv, dropout, l2normalize
from lib.metrics import nnc_score, nnd_score
from lib.theano_utils import floatX, sharedX
from lib.data_utils import OneHot, shuffle, iter_data, center_crop, patch

from scipy.misc import imread
from glob import glob
from load import *

def transform(X):
    return floatX(X)/127.5 - 1

def inverse_transform(X):
    X = (X.reshape(-1, nc, npx, npx).transpose(0, 2, 3, 1)+1.)/2.
    return X

k = 1             # # of discrim updates for each gen update
l2 = 1e-5         # l2 weight decay
nvis = 196        # # of samples to visualize during training
b1 = 0.5          # momentum term of adam
nc = 3            # # of channels in image
nbatch = 100      # # of examples in batch
npx = 32          # # of pixels width/height of images
nz = 100          # # of dim for Z
ngf = 128         # # of gen filters in first conv layer
ndf = 128         # # of discrim filters in first conv layer
ndfc = 1024
nx = npx*npx*nc   # # of dimensions in X
niter = 100       # # of iter at starting learning rate
niter_decay = 100   # # of iter to linearly decay learning rate to zero
ny = 10
margin = 1.

trX, trY, vaX, vaY, teX, teY = load_cifar10()
ntrain, nvalid, ntest = trX.shape[0], vaX.shape[0], teX.shape[0]

desc = 'steingan_supervised'
model_dir = 'models/%s' % desc
samples_dir = 'samples/%s' % desc

dir_list = [model_dir, samples_dir]
for dir in dir_list:
    if not os.path.exists(dir):
        os.makedirs(dir)
print desc

relu = activations.Rectify()
sigmoid = activations.Sigmoid()
lrelu = activations.LeakyRectify()
tanh = activations.Tanh()
bce = T.nnet.binary_crossentropy

gifn = inits.Normal(scale=0.02)
difn = inits.Normal(scale=0.02)
gain_ifn = inits.Normal(loc=1., scale=0.02)
bias_ifn = inits.Constant(c=0.)

gw  = gifn((nz+ny, ngf*4*4*4), 'gw')
gg = gain_ifn((ngf*4*4*4), 'gg')
gb = bias_ifn((ngf*4*4*4), 'gb')
gw2 = gifn((ngf*4+ny, ngf*4, 3, 3), 'gw2')
gg2 = gain_ifn((ngf*4), 'gg2')
gb2 = bias_ifn((ngf*4), 'gb2')
gw3 = gifn((ngf*4+ny, ngf*2, 3, 3), 'gw3')
gg3 = gain_ifn((ngf*2), 'gg3')
gb3 = bias_ifn((ngf*2), 'gb3')
gw4 = gifn((ngf*2+ny, ngf*2, 3, 3), 'gw4')
gg4 = gain_ifn((ngf*2), 'gg4')
gb4 = bias_ifn((ngf*2), 'gb4')
gw5 = gifn((ngf*2+ny, ngf, 3, 3), 'gw5')
gg5 = gain_ifn((ngf), 'gg5')
gb5 = bias_ifn((ngf), 'gb5')
gw6 = gifn((ngf+ny, ngf, 3, 3), 'gw6')
gg6 = gain_ifn((ngf), 'gg6')
gb6 = bias_ifn((ngf), 'gb6')
gwx  = gifn((ngf+ny, nc, 3, 3), 'gwx')

aew1 = difn((ndf, nc, 3, 3), 'aew1')
aew2 = difn((ndf, ndf, 2, 2), 'aew2')
aew3 = difn((ndf*2, ndf, 3, 3), 'aew3')
aew4 = difn((ndf*2, ndf*2, 4, 4), 'aew4')
aew5 = difn((ndf*4, ndf*2, 3, 3), 'aew5')
aew6 = difn((ndf*4, ndf*4, 4, 4), 'aew6')

aeg2 = gain_ifn((ndf), 'aeg2')
aeb2 = bias_ifn((ndf), 'aeb2')
aeg3 = gain_ifn((ndf*2), 'aeg3')
aeb3 = bias_ifn((ndf*2), 'aeb3')
aeg4 = gain_ifn((ndf*2), 'aeg4')
aeb4 = bias_ifn((ndf*2), 'aeb4')
aeg5 = gain_ifn((ndf*4), 'aeg5')
aeb5 = bias_ifn((ndf*4), 'aeb5')
aeg6 = gain_ifn((ndf*4), 'aeg6')
aeb6 = bias_ifn((ndf*4), 'aeb6')

aeg6t = gain_ifn((ndf*4), 'aeg6t')
aeb6t = bias_ifn((ndf*4), 'aeb6t')
aeg5t = gain_ifn((ndf*2), 'aeg5t')
aeb5t = bias_ifn((ndf*2), 'aeb5t')
aeg4t = gain_ifn((ndf*2), 'aeg4t')
aeb4t = bias_ifn((ndf*2), 'aeb4t')
aeg3t = gain_ifn((ndf), 'aeg3t')
aeb3t = bias_ifn((ndf), 'aeb3t')
aeg2t = gain_ifn((ndf), 'aeg2t')
aeb2t = bias_ifn((ndf), 'aeb2t')

logistic_w = difn((ndf*4,ny), 'logistic_w')
logistic_b = bias_ifn((ny,), 'logistic_b')

gen_params = [gw, gg, gb, gw2, gg2, gb2, gw3, gg3, gb3, gw4, gg4, gb4, gw5, gg5, gb5, gw6, gg6, gb6, gwx]
discrim_params = [aew1, aew2, aew3, aew4, aew5, aew6, aeg2, aeb2, aeg3, aeb3, aeg4, aeb4, aeg5, aeb5, aeg6, aeb6, aeg2t, aeb2t, aeg3t, aeb3t, aeg4t, aeb4t, aeg5t, aeb5t, aeg6t, aeb6t, logistic_w, logistic_b]

def gen(Z, Y):
    yb = Y.dimshuffle(0, 1, 'x', 'x')
    Z = T.concatenate([Z, Y], axis=1)
    h = relu(batchnorm(T.dot(Z, gw), g=gg, b=gb))
    h = h.reshape((h.shape[0], ngf*4, 4, 4))
    h = conv_cond_concat(h, yb)
    h2 = relu(batchnorm(deconv(h, gw2, subsample=(2, 2), border_mode=(1, 1)), g=gg2, b=gb2))
    h2 = conv_cond_concat(h2, yb)
    h3 = relu(batchnorm(deconv(h2, gw3, subsample=(1, 1), border_mode=(1, 1)), g=gg3, b=gb3))
    h3 = conv_cond_concat(h3, yb)
    h4 = relu(batchnorm(deconv(h3, gw4, subsample=(2, 2), border_mode=(1, 1)), g=gg4, b=gb4))
    h4 = conv_cond_concat(h4, yb)
    h5 = relu(batchnorm(deconv(h4, gw5, subsample=(1, 1), border_mode=(1, 1)), g=gg5, b=gb5))
    h5 = conv_cond_concat(h5, yb)
    h6 = relu(batchnorm(deconv(h5, gw6, subsample=(2, 2), border_mode=(1, 1)), g=gg6, b=gb6))
    h6 = conv_cond_concat(h6, yb)
    x = tanh(deconv(h6, gwx, subsample=(1, 1), border_mode=(1, 1)))

    return x

def discrim(X, Y):

    def classifier(H, Y):
        p_y_given_x = T.nnet.softmax(T.dot(H, logistic_w) + logistic_b)
        neg_lik = -T.sum(T.mul(T.log(p_y_given_x), Y), axis=1)
        return neg_lik, p_y_given_x

    current_input = dropout(X, 0.2)
    ### encoder ###
    cv1 = relu(dnn_conv(current_input, aew1, subsample=(1,1), border_mode=(1,1)))
    cv2 = relu(batchnorm(dnn_conv(cv1, aew2, subsample=(2,2), border_mode=(0,0)), g=aeg2, b=aeb2))
    cv3 = relu(batchnorm(dnn_conv(cv2, aew3, subsample=(1,1), border_mode=(1,1)), g=aeg3, b=aeb3))
    cv4 = relu(batchnorm(dnn_conv(cv3, aew4, subsample=(4,4), border_mode=(0,0)), g=aeg4, b=aeb4))
    cv5 = relu(batchnorm(dnn_conv(cv4, aew5, subsample=(1,1), border_mode=(1,1)), g=aeg5, b=aeb5))
    cv6 = relu(batchnorm(dnn_conv(cv5, aew6, subsample=(4,4), border_mode=(0,0)), g=aeg6, b=aeb6))

    ### decoder ###
    dv6 = relu(batchnorm(deconv(cv6, aew6, subsample=(4,4), border_mode=(0,0)), g=aeg6t, b=aeb6t))
    dv5 = relu(batchnorm(deconv(dv6, aew5, subsample=(1,1), border_mode=(1,1)), g=aeg5t, b=aeb5t))
    dv4 = relu(batchnorm(deconv(dv5, aew4, subsample=(4,4), border_mode=(0,0)), g=aeg4t, b=aeb4t))
    dv3 = relu(batchnorm(deconv(dv4, aew3, subsample=(1,1), border_mode=(1,1)), g=aeg3t, b=aeb3t))
    dv2 = relu(batchnorm(deconv(dv3, aew2, subsample=(2,2), border_mode=(0,0)), g=aeg2t, b=aeb2t))
    dv1 = tanh(deconv(dv2, aew1, subsample=(1,1), border_mode=(1,1)))

    hidden = T.flatten(cv6, 2)
    rX = dv1
    mse = T.sqrt(T.sum(T.flatten((X-rX)**2, 2), axis=1))

    #mse = T.sqrt(T.sum(T.abs_(T.flatten(X-rX, 2)),axis=1)) + T.sqrt(T.sum(T.flatten((X-rX)**2, 2), axis=1))
    neg_lik, p_y_given_x = classifier(hidden, Y)
    return hidden, p_y_given_x, rX, mse, neg_lik

def combine_err(mse, neg_lik):

    return mse + T.maximum(margin, neg_lik)

def rbf_kernel(X0):
    xy = T.dot(X0, X0.T)
    x2 = T.sum(X0**2, axis=1).dimshuffle(0, 'x')
    x2e = T.tile(x2, (1, x2.shape[0]))
    H = -2. * xy
    H += x2e
    H += x2e.T

    V = H.flatten()

    # median distance
    h = T.switch(T.eq((V.shape[0] % 2), 0),
        # if even vector
        T.mean(T.sort(V)[ ((V.shape[0] // 2) - 1) : ((V.shape[0] // 2) + 1) ]),
        # if odd vector
        T.sort(V)[V.shape[0] // 2])

    h = T.sqrt(0.5 * h / T.log(X0.shape[0].astype('float32') + 1.0)) / 2.
    Kxy = T.exp(-H / h ** 2 / 2.0)

    neighbors = T.argsort(H, axis=1)[:, 1]
    return Kxy, neighbors, h

def gradient(X0, Y):

    _, _, _, mse, neg_lik = discrim(X0, Y)
    grad = -1.0 * T.grad(combine_err(mse, neg_lik).mean() , X0)

    return grad

def svgd_gradient(X0, Y):

    hidden, _, _, mse, neg_lik = discrim(X0, Y)
    grad = -1.0 * T.grad( combine_err(mse, neg_lik).sum(), X0)

    kxy, neighbors, h = rbf_kernel(hidden)  #TODO

    coff = T.exp( - T.sum((hidden[neighbors] - hidden)**2, axis=1) / h**2 / 2.0 )
    v = coff.dimshuffle(0, 'x') * (-hidden[neighbors] + hidden) / h**2

    X1 = X0[neighbors]
    hidden1, _, _, _, _ = discrim(X1, Y)
    dxkxy = T.Lop(hidden1, X1, v)

    #svgd_grad = (T.dot(kxy, T.flatten(grad, 2)).reshape(dxkxy.shape) + dxkxy) / T.sum(kxy, axis=1).dimshuffle(0, 'x', 'x', 'x')
    svgd_grad = grad + dxkxy / 2.
    return grad, svgd_grad

X = T.tensor4() # data

X0 = T.tensor4() # vgd samples

deltaX = T.tensor4() #vgd gradient

Z = T.matrix()

Y = T.matrix()

epsilon = T.tensor4()

### define discriminative cost ###
H_data, py_data, rX_data, mse_data, neg_lik_data = discrim(X, Y)
H_vgd, py_data, rX_vgd, mse_vgd, neg_lik_vgd = discrim(X0, Y)

cost_data = combine_err(mse_data, neg_lik_data).mean()
cost_vgd = combine_err(mse_vgd, neg_lik_vgd).mean()

balance_weight = sharedX(.3)
d_cost = cost_data - balance_weight * cost_vgd   # for discriminative model, minimize cost

gX = gen(Z, Y)
g_cost = -1 * T.sum(T.sum(T.flatten(gX, 2) * T.flatten(deltaX, 2), axis=1)) #update generate models by minimize reconstruct mse

cost = [mse_data.mean(), neg_lik_data.mean(), mse_vgd.mean(), neg_lik_vgd.mean()]

d_lr = 5e-4
g_lr = 1e-3

d_lrt = sharedX(d_lr)
g_lrt = sharedX(g_lr)

d_updater = updates.Adam(lr=d_lrt, b1=b1, regularizer=updates.Regularizer(l2=l2))
g_updater = updates.Adam(lr=g_lrt, b1=b1, regularizer=updates.Regularizer(l2=l2))

d_updates = d_updater(discrim_params, d_cost)
g_updates = g_updater(gen_params, g_cost)

print 'COMPILING'
t = time()
_gen = theano.function([Z, Y], gX)
_discrim = theano.function([X, Y], discrim(X, Y))
_train_d = theano.function([X, X0, Y], cost, updates=d_updates)
_train_g = theano.function([Z, Y, deltaX], g_cost, updates=g_updates)
_svgd_gradient = theano.function([X0, Y], svgd_gradient(X0, Y))
_gradient = theano.function([X0, Y], gradient(X0, Y))
_reconstruction_cost = theano.function([X], T.mean(mse_data))
_reconstruction = theano.function([X], rX_data)
print '%.2f seconds to compile theano functions'%(time()-t)

sample_zmb = floatX(np_rng.uniform(-1., 1., size=(200, nz)))
sample_ymb = floatX(OneHot(np.asarray([[i for _ in range(20)] for i in range(10)]).flatten(), ny))

n_updates = 0
n_epochs = 0

t = time()
for epoch in range(1, 200+1):

    print 'cifar 10, steigan, %s, iter %d' % (desc, epoch)
    trX, trY = shuffle(trX, trY)

    for imb, ymb in tqdm(iter_data(trX, trY, size=nbatch), total=ntrain/nbatch):
        imb = transform(imb.reshape(imb.shape[0], nc, npx, npx))
        ymb = floatX(OneHot(ymb, ny))
        zmb = floatX(np_rng.uniform(-1., 1., size=(imb.shape[0], nz)))
        # generate samples
        samples = _gen(zmb, ymb)

        # G
        grad, vgd_grad = _svgd_gradient(samples, ymb)
        _train_g(zmb, ymb, floatX(vgd_grad))

        # D
        cost = _train_d(imb, samples, ymb)

        n_updates += 1

        cost_batch_vgd = _reconstruction_cost(floatX(samples))
        cost_batch_data = _reconstruction_cost(imb)

        if n_updates % 10 == 0:
            print desc, cost_batch_data, cost_batch_vgd

        if cost_batch_data > cost_batch_vgd:
            balance_weight.set_value(0.3)
        else:
            balance_weight.set_value(0.1)

    samples = np.asarray(_gen(sample_zmb, sample_ymb))
    color_grid_vis(inverse_transform(samples), (10, 20), 'samples/%s/gan-%d.png' % (desc, epoch))
    color_grid_vis(inverse_transform(_reconstruction(imb)), (10, 10), 'samples/%s/ae-%d.png' % (desc, epoch))

    n_epochs += 1

    if epoch % 50 == 0:
        joblib.dump([p.get_value() for p in gen_params], 'models/%s/%d_gen_params.jl'%(desc, epoch))
        joblib.dump([p.get_value() for p in discrim_params], 'models/%s/%d_discrim_params.jl'%(desc, epoch))

print '%.2f seconds to train the generative model' % (time()-t)
print 'DONE'

<commit_msg>Delete steingan_supervised.py<commit_after><|endoftext|>