<commit_before>from __future__ import print_function, unicode_literals

import argparse
import sys
import os
import os.path as path
import subprocess
from time import time

from mach.registrar import Registrar
from mach.decorators import (
    CommandArgument,
    CommandProvider,
    Command,
)

from servo.command_base import CommandBase
import tidy
import multiprocessing

@CommandProvider
class MachCommands(CommandBase):
    def __init__(self, context):
        CommandBase.__init__(self, context)
        if not hasattr(self.context, "built_tests"):
            self.context.built_tests = False

    def ensure_built_tests(self):
        if self.context.built_tests:
            return
        returncode = Registrar.dispatch('build-tests', context=self.context)
	if returncode:
	    sys.exit(returncode)
        self.context.built_tests = True

    def find_test(self, prefix):
        target_contents = os.listdir(path.join(
            self.context.topdir, "components", "servo", "target"))
        for filename in target_contents:
            if filename.startswith(prefix + "-"):
                filepath = path.join(
                    self.context.topdir, "components", "servo",
                    "target", filename)

                if path.isfile(filepath) and os.access(filepath, os.X_OK):
                    return filepath

    def run_test(self, prefix, args=[]):
        t = self.find_test(prefix)
        if t:
            return subprocess.call([t] + args, env=self.build_env())

    def infer_test_by_dir(self, params):
        maybe_path = path.normpath(params[0])
        mach_command = path.join(self.context.topdir, "mach")
        args = None

        if not path.exists(maybe_path):
            print("%s is not a valid file or directory" % maybe_path)
            return 1

        test_dirs = [
            # path, mach test command, optional flag for path argument
            (path.join("tests", "content"), "test-content", None),
            (path.join("tests", "wpt"), "test-wpt", None),
            (path.join("tests", "ref"), "test-ref", ["--name"]),
        ]

        for test_dir, test_name, path_flag in test_dirs:
            if not path_flag:
                path_flag = []
            if test_dir in maybe_path:
                args = ([mach_command, test_name] + path_flag +
                        [maybe_path] + params[1:])
                break
        else:
            print("%s is not a valid test file or directory" % maybe_path)
            return 1

        return subprocess.call(args, env=self.build_env())

    @Command('test',
             description='Run all Servo tests',
             category='testing')
    @CommandArgument('params', default=None, nargs="...",
                     help="Optionally select test based on "
                          "test file directory")
    def test(self, params):
        if params:
            return self.infer_test_by_dir(params)

        test_start = time()
        for t in ["tidy", "ref", "content", "wpt", "unit"]:
            Registrar.dispatch("test-%s" % t, context=self.context)
        elapsed = time() - test_start

        print("Tests completed in %0.2fs" % elapsed)

    @Command('test-unit',
             description='Run unit tests',
             category='testing')
    @CommandArgument('--component', '-c', default=None,
                     help="Specific component to test")
    @CommandArgument('test_name', nargs=argparse.REMAINDER,
                     help="Only run tests that match this pattern")
    def test_unit(self, test_name=None, component=None):
        if test_name is None:
            test_name = []

        self.ensure_bootstrapped()
        self.ensure_built_tests()

        ret = self.run_test("servo", test_name) != 0

        def cargo_test(component):
            return 0 != subprocess.call(
                ["cargo", "test", "-p", component]
                + test_name, env=self.build_env(), cwd=self.servo_crate())

        if component is not None:
            ret = ret or cargo_test(component)
        else:
            for c in os.listdir("components"):
                if c != "servo":
                    ret = ret or cargo_test(c)

        return ret

    @Command('test-ref',
             description='Run the reference tests',
             category='testing')
    @CommandArgument('--kind', '-k', default=None,
                     help="'cpu' or 'gpu' (default both)")
    @CommandArgument('--name', default=None,
                     help="Only run tests that match this pattern. If the "
                          "path to the ref test directory is included, it "
                          "will automatically be trimmed out.")
    @CommandArgument(
        'servo_params', default=None, nargs=argparse.REMAINDER,
        help="Command-line arguments to be passed through to Servo")
    def test_ref(self, kind=None, name=None, servo_params=None):
        self.ensure_bootstrapped()
        self.ensure_built_tests()

        kinds = ["cpu", "gpu"] if kind is None else [kind]
        test_path = path.join(self.context.topdir, "tests", "ref")
        error = False

        test_start = time()
        for k in kinds:
            print("Running %s reftests..." % k)
            test_args = [k, test_path]
            if name is not None:
                maybe_path = path.normpath(name)
                ref_path = path.join("tests", "ref")

                # Check to see if we were passed something leading with the
                # path to the ref test directory, and trim it so that reftest
                # knows how to filter it.
                if ref_path in maybe_path:
                    test_args.append(path.relpath(maybe_path, ref_path))
                else:
                    test_args.append(name)
            if servo_params is not None:
                test_args += ["--"] + servo_params
            ret = self.run_test("reftest", test_args)
            error = error or ret != 0
        elapsed = time() - test_start

        print("Reference tests completed in %0.2fs" % elapsed)

        if error:
            return 1

    @Command('test-content',
             description='Run the content tests',
             category='testing')
    @CommandArgument('test_name', default=None, nargs="?",
                     help="Only run tests that match this pattern")
    def test_content(self, test_name=None):
        self.ensure_bootstrapped()
        self.ensure_built_tests()
        test_path = path.join(self.context.topdir, "tests", "content")
        test_args = ["--source-dir=%s" % test_path]

        if test_name is not None:
            test_args.append(test_name)

        test_start = time()
        ret = self.run_test("contenttest", test_args)
        elapsed = time() - test_start

        print("Content tests completed in %0.2fs" % elapsed)
        return ret

    @Command('test-tidy',
             description='Run the source code tidiness check',
             category='testing')
    def test_tidy(self):
        return tidy.scan()

    @Command('test-wpt-failure',
             description='Run the web platform tests',
             category='testing')
    def test_wpt_failure(self):
        return not subprocess.call([
            "bash",
            path.join("tests", "wpt", "run.sh"),
            "--include",
            "infrastructure/failing-test.html"
        ], env=self.build_env())

    @Command('test-wpt',
             description='Run the web platform tests',
             category='testing')
    @CommandArgument(
        '--processes', default=None,
        help="Number of servo processes to spawn")
    @CommandArgument(
        "params", default=None, nargs='...',
        help="command-line arguments to be passed through to wpt/run.sh")
    def test_wpt(self, processes=None, params=None):
        if params is None:
            params = []
        else:
            # Allow the first argument to be a relative path from Servo's root
            # directory, converting it to `--include some/wpt/test.html`
            maybe_path = path.normpath(params[0])
            wpt_path = path.join('tests', 'wpt', 'web-platform-tests')

            if path.exists(maybe_path) and wpt_path in maybe_path:
                params = ["--include",
                          path.relpath(maybe_path, wpt_path)]

        processes = str(multiprocessing.cpu_count()) if processes is None else processes
        params = params + ["--processes", processes]

        return subprocess.call(
            ["bash", path.join("tests", "wpt", "run.sh")] + params,
            env=self.build_env())
<commit_msg>tests: ensure bootstrapping for wpt<commit_after>from __future__ import print_function, unicode_literals

import argparse
import sys
import os
import os.path as path
import subprocess
from time import time

from mach.registrar import Registrar
from mach.decorators import (
    CommandArgument,
    CommandProvider,
    Command,
)

from servo.command_base import CommandBase
import tidy
import multiprocessing

@CommandProvider
class MachCommands(CommandBase):
    def __init__(self, context):
        CommandBase.__init__(self, context)
        if not hasattr(self.context, "built_tests"):
            self.context.built_tests = False

    def ensure_built_tests(self):
        if self.context.built_tests:
            return
        returncode = Registrar.dispatch('build-tests', context=self.context)
	if returncode:
	    sys.exit(returncode)
        self.context.built_tests = True

    def find_test(self, prefix):
        target_contents = os.listdir(path.join(
            self.context.topdir, "components", "servo", "target"))
        for filename in target_contents:
            if filename.startswith(prefix + "-"):
                filepath = path.join(
                    self.context.topdir, "components", "servo",
                    "target", filename)

                if path.isfile(filepath) and os.access(filepath, os.X_OK):
                    return filepath

    def run_test(self, prefix, args=[]):
        t = self.find_test(prefix)
        if t:
            return subprocess.call([t] + args, env=self.build_env())

    def infer_test_by_dir(self, params):
        maybe_path = path.normpath(params[0])
        mach_command = path.join(self.context.topdir, "mach")
        args = None

        if not path.exists(maybe_path):
            print("%s is not a valid file or directory" % maybe_path)
            return 1

        test_dirs = [
            # path, mach test command, optional flag for path argument
            (path.join("tests", "content"), "test-content", None),
            (path.join("tests", "wpt"), "test-wpt", None),
            (path.join("tests", "ref"), "test-ref", ["--name"]),
        ]

        for test_dir, test_name, path_flag in test_dirs:
            if not path_flag:
                path_flag = []
            if test_dir in maybe_path:
                args = ([mach_command, test_name] + path_flag +
                        [maybe_path] + params[1:])
                break
        else:
            print("%s is not a valid test file or directory" % maybe_path)
            return 1

        return subprocess.call(args, env=self.build_env())

    @Command('test',
             description='Run all Servo tests',
             category='testing')
    @CommandArgument('params', default=None, nargs="...",
                     help="Optionally select test based on "
                          "test file directory")
    def test(self, params):
        if params:
            return self.infer_test_by_dir(params)

        test_start = time()
        for t in ["tidy", "ref", "content", "wpt", "unit"]:
            Registrar.dispatch("test-%s" % t, context=self.context)
        elapsed = time() - test_start

        print("Tests completed in %0.2fs" % elapsed)

    @Command('test-unit',
             description='Run unit tests',
             category='testing')
    @CommandArgument('--component', '-c', default=None,
                     help="Specific component to test")
    @CommandArgument('test_name', nargs=argparse.REMAINDER,
                     help="Only run tests that match this pattern")
    def test_unit(self, test_name=None, component=None):
        if test_name is None:
            test_name = []

        self.ensure_bootstrapped()
        self.ensure_built_tests()

        ret = self.run_test("servo", test_name) != 0

        def cargo_test(component):
            return 0 != subprocess.call(
                ["cargo", "test", "-p", component]
                + test_name, env=self.build_env(), cwd=self.servo_crate())

        if component is not None:
            ret = ret or cargo_test(component)
        else:
            for c in os.listdir("components"):
                if c != "servo":
                    ret = ret or cargo_test(c)

        return ret

    @Command('test-ref',
             description='Run the reference tests',
             category='testing')
    @CommandArgument('--kind', '-k', default=None,
                     help="'cpu' or 'gpu' (default both)")
    @CommandArgument('--name', default=None,
                     help="Only run tests that match this pattern. If the "
                          "path to the ref test directory is included, it "
                          "will automatically be trimmed out.")
    @CommandArgument(
        'servo_params', default=None, nargs=argparse.REMAINDER,
        help="Command-line arguments to be passed through to Servo")
    def test_ref(self, kind=None, name=None, servo_params=None):
        self.ensure_bootstrapped()
        self.ensure_built_tests()

        kinds = ["cpu", "gpu"] if kind is None else [kind]
        test_path = path.join(self.context.topdir, "tests", "ref")
        error = False

        test_start = time()
        for k in kinds:
            print("Running %s reftests..." % k)
            test_args = [k, test_path]
            if name is not None:
                maybe_path = path.normpath(name)
                ref_path = path.join("tests", "ref")

                # Check to see if we were passed something leading with the
                # path to the ref test directory, and trim it so that reftest
                # knows how to filter it.
                if ref_path in maybe_path:
                    test_args.append(path.relpath(maybe_path, ref_path))
                else:
                    test_args.append(name)
            if servo_params is not None:
                test_args += ["--"] + servo_params
            ret = self.run_test("reftest", test_args)
            error = error or ret != 0
        elapsed = time() - test_start

        print("Reference tests completed in %0.2fs" % elapsed)

        if error:
            return 1

    @Command('test-content',
             description='Run the content tests',
             category='testing')
    @CommandArgument('test_name', default=None, nargs="?",
                     help="Only run tests that match this pattern")
    def test_content(self, test_name=None):
        self.ensure_bootstrapped()
        self.ensure_built_tests()
        test_path = path.join(self.context.topdir, "tests", "content")
        test_args = ["--source-dir=%s" % test_path]

        if test_name is not None:
            test_args.append(test_name)

        test_start = time()
        ret = self.run_test("contenttest", test_args)
        elapsed = time() - test_start

        print("Content tests completed in %0.2fs" % elapsed)
        return ret

    @Command('test-tidy',
             description='Run the source code tidiness check',
             category='testing')
    def test_tidy(self):
        return tidy.scan()

    @Command('test-wpt-failure',
             description='Run the web platform tests',
             category='testing')
    def test_wpt_failure(self):
        self.ensure_bootstrapped()
        return not subprocess.call([
            "bash",
            path.join("tests", "wpt", "run.sh"),
            "--include",
            "infrastructure/failing-test.html"
        ], env=self.build_env())

    @Command('test-wpt',
             description='Run the web platform tests',
             category='testing')
    @CommandArgument(
        '--processes', default=None,
        help="Number of servo processes to spawn")
    @CommandArgument(
        "params", default=None, nargs='...',
        help="command-line arguments to be passed through to wpt/run.sh")
    def test_wpt(self, processes=None, params=None):
        self.ensure_bootstrapped()

        if params is None:
            params = []
        else:
            # Allow the first argument to be a relative path from Servo's root
            # directory, converting it to `--include some/wpt/test.html`
            maybe_path = path.normpath(params[0])
            wpt_path = path.join('tests', 'wpt', 'web-platform-tests')

            if path.exists(maybe_path) and wpt_path in maybe_path:
                params = ["--include",
                          path.relpath(maybe_path, wpt_path)]

        processes = str(multiprocessing.cpu_count()) if processes is None else processes
        params = params + ["--processes", processes]

        return subprocess.call(
            ["bash", path.join("tests", "wpt", "run.sh")] + params,
            env=self.build_env())
<|endoftext|>


<commit_before>import json
import logging
from collections import defaultdict

import elasticsearch
import elasticsearch_dsl
from elasticsearch_async.transport import AsyncTransport
from elasticsearch_dsl.connections import Connections
from tornado.ioloop import IOLoop

from biothings.utils.web.es import get_es_versions
from biothings.utils.web.userquery import ESUserQuery
from biothings.web.settings import BiothingWebSettings

class BiothingESWebSettings(BiothingWebSettings):
    '''
    `BiothingWebSettings`_ subclass with functions specific to an elasticsearch backend.

    * Use the known live ES connection if more than one is specified.
    * Cache source metadata stored under the _meta field in es indices.

    '''

    def __init__(self, config=None, **kwargs):
        '''
        The ``config`` init parameter specifies a module that configures
        this biothing.  For more information see `config module`_ documentation.
        '''
        super(BiothingESWebSettings, self).__init__(config, **kwargs)

        # elasticsearch connections
        self._connections = Connections()

        connection_settings = {
            "hosts": self.ES_HOST,
            "timeout": self.ES_CLIENT_TIMEOUT,
            "max_retries": 1,  # maximum number of retries before an exception is propagated
            "timeout_cutoff": 1,  # timeout freezes after this number of consecutive failures
            # "sniff_on_connection_fail": True, TODO
            "sniffer_timeout": 60
        }
        self._connections.create_connection(alias='sync', **connection_settings)
        connection_settings.update(transport_class=AsyncTransport)
        self._connections.create_connection(alias='async', **connection_settings)

        # cached index mappings
        self.source_metadata = defaultdict(dict)
        self.source_properties = defaultdict(dict)

        # populate field notes if exist
        try:
            inf = open(self.AVAILABLE_FIELDS_NOTES_PATH, 'r')
            self._fields_notes = json.load(inf)
            inf.close()
        except Exception:
            self._fields_notes = {}

        # user query data
        self.userquery = ESUserQuery(self.USERQUERY_DIR)

        # query pipelines
        self.query_builder = self.load_class(self.ES_QUERY_BUILDER)(self)
        self.query_backend = self.load_class(self.ES_QUERY_BACKEND)(self)
        self.result_transform = self.load_class(self.ES_RESULT_TRANSFORM)(self)

        # initialize payload for standalone tracking batch
        self.tracking_payload = []

        self.ES_INDICES = dict(self.ES_INDICES)  # TODO
        self.ES_INDICES[self.ES_DOC_TYPE] = self.ES_INDEX
        self.BIOTHING_TYPES = list(self.ES_INDICES.keys())

        IOLoop.current().add_callback(self._initialize)

    async def _initialize(self):

        # failures will be logged concisely
        logging.getLogger('elasticsearch.trace').propagate = False

        self.logger.info("Python Elasticsearch Version: %s", elasticsearch.__version__)
        self.logger.info("Python Elasticsearch DSL Version: %s", elasticsearch_dsl.__version__)
        versions = await get_es_versions(self.async_es_client)
        versions = iter(versions.values())
        self.logger.info('Elasticsearch Version: %s', next(versions))
        self.logger.info('Elasticsearch Cluster: %s', next(versions))

        # populate source mappings
        for biothing_type in self.ES_INDICES:
            await self.read_index_mappings(biothing_type)

        # resume normal log flow
        logging.getLogger('elasticsearch.trace').propagate = True

    def validate(self):
        '''
        Additional ES settings to validate.
        '''
        super().validate()

        assert isinstance(self.ES_INDEX, str)
        assert isinstance(self.ES_DOC_TYPE, str)
        assert isinstance(self.ES_INDICES, dict)
        assert '*' not in self.ES_DOC_TYPE

    def get_es_client(self):
        '''
        Return the default blocking elasticsearch client.
        The connection is created upon first call.
        '''
        return self._connections.get_connection('sync')

    def get_async_es_client(self):
        '''
        Return the async elasitcsearch client.
        The connection is created upon first call.
        API calls return awaitable objects.
        '''
        return self._connections.get_connection('async')

    async def read_index_mappings(self, biothing_type=None):
        """
        Read ES index mappings for the corresponding biothing_type,
        Populate datasource info and field properties from mappings.
        Return ES raw response. This implementation combines indices.

        The ES response would look like: (for es7+)
        {
            'index_1': {
                'properties': { ... },  ---> source_properties
                '_meta': {
                    "src" : { ... }     ---> source_licenses
                    ...
                },              -----------> source_metadata
                ...
            },
            'index_2': {
                ...     ---------> Combine with results above
            }
        }
        """

        biothing_type = biothing_type or self.ES_DOC_TYPE
        try:
            mappings = await self.async_es_client.indices.get_mapping(
                index=self.ES_INDICES[biothing_type],
                allow_no_indices=True,
                ignore_unavailable=True,
                local=False)
        except elasticsearch.TransportError as exc:
            self.logger.error('Error loading index mapping for [%s].', biothing_type)
            self.logger.debug(exc)
            return None

        metadata = self.source_metadata[biothing_type]
        properties = self.source_properties[biothing_type]
        licenses = self.result_transform.source_licenses[biothing_type]

        metadata.clear()
        properties.clear()
        licenses.clear()

        for index in mappings:

            mapping = mappings[index]['mappings']

            if elasticsearch.__version__[0] < 7:
                # remove doc_type, support 1 type per index
                mapping = next(iter(mapping.values()))

            if '_meta' in mapping:
                for key, val in mapping['_meta'].items():
                    if key in metadata and isinstance(val, dict) \
                            and isinstance(metadata[key], dict):
                        metadata[key].update(val)
                    else:
                        metadata[key] = val
                metadata.update(mapping['_meta'])

                if 'src' in mapping['_meta']:
                    for src, info in mapping['_meta']['src'].items():
                        if 'license_url_short' in info:
                            licenses[src] = info['license_url_short']
                        elif 'license_url' in info:
                            licenses[src] = info['license_url']

            if 'properties' in mapping:
                properties.update(mapping['properties'])

        return mappings

    def get_field_notes(self):
        '''
        Return the cached field notes associated with this instance.
        '''
        return self._fields_notes

    @property
    def es_client(self):
        return self.get_es_client()

    @property
    def async_es_client(self):
        return self.get_async_es_client()
<commit_msg>Disable sniffing for now<commit_after>import json
import logging
from collections import defaultdict

import elasticsearch
import elasticsearch_dsl
from elasticsearch_async.transport import AsyncTransport
from elasticsearch_dsl.connections import Connections
from tornado.ioloop import IOLoop

from biothings.utils.web.es import get_es_versions
from biothings.utils.web.userquery import ESUserQuery
from biothings.web.settings import BiothingWebSettings

class BiothingESWebSettings(BiothingWebSettings):
    '''
    `BiothingWebSettings`_ subclass with functions specific to an elasticsearch backend.

    * Use the known live ES connection if more than one is specified.
    * Cache source metadata stored under the _meta field in es indices.

    '''

    def __init__(self, config=None, **kwargs):
        '''
        The ``config`` init parameter specifies a module that configures
        this biothing.  For more information see `config module`_ documentation.
        '''
        super(BiothingESWebSettings, self).__init__(config, **kwargs)

        # elasticsearch connections
        self._connections = Connections()

        connection_settings = {
            "hosts": self.ES_HOST,
            "timeout": self.ES_CLIENT_TIMEOUT,
            "max_retries": 1,  # maximum number of retries before an exception is propagated
            "timeout_cutoff": 1,  # timeout freezes after this number of consecutive failures
            # "sniff_on_connection_fail": True, TODO
            # "sniffer_timeout": 60
        }
        self._connections.create_connection(alias='sync', **connection_settings)
        connection_settings.update(transport_class=AsyncTransport)
        self._connections.create_connection(alias='async', **connection_settings)

        # cached index mappings
        self.source_metadata = defaultdict(dict)
        self.source_properties = defaultdict(dict)

        # populate field notes if exist
        try:
            inf = open(self.AVAILABLE_FIELDS_NOTES_PATH, 'r')
            self._fields_notes = json.load(inf)
            inf.close()
        except Exception:
            self._fields_notes = {}

        # user query data
        self.userquery = ESUserQuery(self.USERQUERY_DIR)

        # query pipelines
        self.query_builder = self.load_class(self.ES_QUERY_BUILDER)(self)
        self.query_backend = self.load_class(self.ES_QUERY_BACKEND)(self)
        self.result_transform = self.load_class(self.ES_RESULT_TRANSFORM)(self)

        # initialize payload for standalone tracking batch
        self.tracking_payload = []

        self.ES_INDICES = dict(self.ES_INDICES)  # TODO
        self.ES_INDICES[self.ES_DOC_TYPE] = self.ES_INDEX
        self.BIOTHING_TYPES = list(self.ES_INDICES.keys())

        IOLoop.current().add_callback(self._initialize)

    async def _initialize(self):

        # failures will be logged concisely
        logging.getLogger('elasticsearch.trace').propagate = False

        self.logger.info("Python Elasticsearch Version: %s", elasticsearch.__version__)
        self.logger.info("Python Elasticsearch DSL Version: %s", elasticsearch_dsl.__version__)
        versions = await get_es_versions(self.async_es_client)
        versions = iter(versions.values())
        self.logger.info('Elasticsearch Version: %s', next(versions))
        self.logger.info('Elasticsearch Cluster: %s', next(versions))

        # populate source mappings
        for biothing_type in self.ES_INDICES:
            await self.read_index_mappings(biothing_type)

        # resume normal log flow
        logging.getLogger('elasticsearch.trace').propagate = True

    def validate(self):
        '''
        Additional ES settings to validate.
        '''
        super().validate()

        assert isinstance(self.ES_INDEX, str)
        assert isinstance(self.ES_DOC_TYPE, str)
        assert isinstance(self.ES_INDICES, dict)
        assert '*' not in self.ES_DOC_TYPE

    def get_es_client(self):
        '''
        Return the default blocking elasticsearch client.
        The connection is created upon first call.
        '''
        return self._connections.get_connection('sync')

    def get_async_es_client(self):
        '''
        Return the async elasitcsearch client.
        The connection is created upon first call.
        API calls return awaitable objects.
        '''
        return self._connections.get_connection('async')

    async def read_index_mappings(self, biothing_type=None):
        """
        Read ES index mappings for the corresponding biothing_type,
        Populate datasource info and field properties from mappings.
        Return ES raw response. This implementation combines indices.

        The ES response would look like: (for es7+)
        {
            'index_1': {
                'properties': { ... },  ---> source_properties
                '_meta': {
                    "src" : { ... }     ---> source_licenses
                    ...
                },              -----------> source_metadata
                ...
            },
            'index_2': {
                ...     ---------> Combine with results above
            }
        }
        """

        biothing_type = biothing_type or self.ES_DOC_TYPE
        try:
            mappings = await self.async_es_client.indices.get_mapping(
                index=self.ES_INDICES[biothing_type],
                allow_no_indices=True,
                ignore_unavailable=True,
                local=False)
        except elasticsearch.TransportError as exc:
            self.logger.error('Error loading index mapping for [%s].', biothing_type)
            self.logger.debug(exc)
            return None

        metadata = self.source_metadata[biothing_type]
        properties = self.source_properties[biothing_type]
        licenses = self.result_transform.source_licenses[biothing_type]

        metadata.clear()
        properties.clear()
        licenses.clear()

        for index in mappings:

            mapping = mappings[index]['mappings']

            if elasticsearch.__version__[0] < 7:
                # remove doc_type, support 1 type per index
                mapping = next(iter(mapping.values()))

            if '_meta' in mapping:
                for key, val in mapping['_meta'].items():
                    if key in metadata and isinstance(val, dict) \
                            and isinstance(metadata[key], dict):
                        metadata[key].update(val)
                    else:
                        metadata[key] = val
                metadata.update(mapping['_meta'])

                if 'src' in mapping['_meta']:
                    for src, info in mapping['_meta']['src'].items():
                        if 'license_url_short' in info:
                            licenses[src] = info['license_url_short']
                        elif 'license_url' in info:
                            licenses[src] = info['license_url']

            if 'properties' in mapping:
                properties.update(mapping['properties'])

        return mappings

    def get_field_notes(self):
        '''
        Return the cached field notes associated with this instance.
        '''
        return self._fields_notes

    @property
    def es_client(self):
        return self.get_es_client()

    @property
    def async_es_client(self):
        return self.get_async_es_client()
<|endoftext|>


<commit_before>""" Joblib is a set of tools to provide **lightweight pipelining in
Python**. In particular, joblib offers:

  1. transparent disk-caching of the output values and lazy re-evaluation
     (memoize pattern)

  2. easy simple parallel computing

  3. logging and tracing of the execution

Joblib is optimized to be **fast** and **robust** in particular on large
data and has specific optimizations for `numpy` arrays. It is
**BSD-licensed**.

    ============================== ============================================
    **User documentation**:        http://packages.python.org/joblib

    **Download packages**:         http://pypi.python.org/pypi/joblib#downloads

    **Source code**:               http://github.com/joblib/joblib

    **Report issues**:             http://github.com/joblib/joblib/issues
    ============================== ============================================

Vision
--------

The vision is to provide tools to easily achieve better performance and
reproducibility when working with long running jobs. In addition, Joblib
can also be used to provide a light-weight make replacement or caching
solution.

 *  **Avoid computing twice the same thing**: code is rerun over an
    over, for instance when prototyping computational-heavy jobs (as in
    scientific development), but hand-crafted solution to aleviate this
    issue is error-prone and often leads to unreproducible results

 *  **Persist to disk transparently**: persisting in an efficient way
    arbitrary objects containing large data is hard. In addition,
    hand-written persistence does not link easily the file on disk to the
    execution context of the original Python object. As a result, it is
    challenging to resume a application status or computational job, eg
    after a crash.

It strives to address these problems while **leaving your code and your
flow control as unmodified as possible** (no framework, no new
paradigms).

Main features
------------------

1) **Transparent and fast disk-caching of output value:** a memoize or
   make-like functionality for Python functions that works well for
   arbitrary Python objects, including very large numpy arrays. Separate
   persistence and flow-execution logic from domain logic or algorithmic
   code by writing the operations as a set of steps with well-defined
   inputs and  outputs: Python functions. Joblib can save their
   computation to disk and rerun it only if necessary::

      >>> from joblib import Memory
      >>> mem = Memory(cachedir='/tmp/joblib')
      >>> import numpy as np
      >>> a = np.vander(np.arange(3))
      >>> square = mem.cache(np.square)
      >>> b = square(a)
      ________________________________________________________________________________
      [Memory] Calling square...
      square(array([[0, 0, 1],
             [1, 1, 1],
             [4, 2, 1]]))
      ___________________________________________________________square - 0.0s, 0.0min

      >>> c = square(a)
      >>> # The above call did not trigger an evaluation

2) **Embarrassingly parallel helper:** to make is easy to write readable
   parallel code and debug it quickly:

      >>> from joblib import Parallel, delayed
      >>> from math import sqrt
      >>> Parallel(n_jobs=1)(delayed(sqrt)(i**2) for i in range(10))
      [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]

3) **Logging/tracing:** The different functionalities will
   progressively acquire better logging mechanism to help track what
   has been ran, and capture I/O easily. In addition, Joblib will
   provide a few I/O primitives, to easily define define logging and
   display streams, and provide a way of compiling a report.
   We want to be able to quickly inspect what has been run.

..
    >>> import shutil ; shutil.rmtree('/tmp/joblib/')

"""

__version__ = '0.5.3'

from .memory import Memory
from .logger import PrintTime, Logger
from .hashing import hash
from .numpy_pickle import dump, load
from .parallel import Parallel, delayed, cpu_count
<commit_msg>FIX: doctest heisenfailure on execution time<commit_after>""" Joblib is a set of tools to provide **lightweight pipelining in
Python**. In particular, joblib offers:

  1. transparent disk-caching of the output values and lazy re-evaluation
     (memoize pattern)

  2. easy simple parallel computing

  3. logging and tracing of the execution

Joblib is optimized to be **fast** and **robust** in particular on large
data and has specific optimizations for `numpy` arrays. It is
**BSD-licensed**.

    ============================== ============================================
    **User documentation**:        http://packages.python.org/joblib

    **Download packages**:         http://pypi.python.org/pypi/joblib#downloads

    **Source code**:               http://github.com/joblib/joblib

    **Report issues**:             http://github.com/joblib/joblib/issues
    ============================== ============================================

Vision
--------

The vision is to provide tools to easily achieve better performance and
reproducibility when working with long running jobs. In addition, Joblib
can also be used to provide a light-weight make replacement or caching
solution.

 *  **Avoid computing twice the same thing**: code is rerun over an
    over, for instance when prototyping computational-heavy jobs (as in
    scientific development), but hand-crafted solution to aleviate this
    issue is error-prone and often leads to unreproducible results

 *  **Persist to disk transparently**: persisting in an efficient way
    arbitrary objects containing large data is hard. In addition,
    hand-written persistence does not link easily the file on disk to the
    execution context of the original Python object. As a result, it is
    challenging to resume a application status or computational job, eg
    after a crash.

It strives to address these problems while **leaving your code and your
flow control as unmodified as possible** (no framework, no new
paradigms).

Main features
------------------

1) **Transparent and fast disk-caching of output value:** a memoize or
   make-like functionality for Python functions that works well for
   arbitrary Python objects, including very large numpy arrays. Separate
   persistence and flow-execution logic from domain logic or algorithmic
   code by writing the operations as a set of steps with well-defined
   inputs and  outputs: Python functions. Joblib can save their
   computation to disk and rerun it only if necessary::

      >>> from joblib import Memory
      >>> mem = Memory(cachedir='/tmp/joblib')
      >>> import numpy as np
      >>> a = np.vander(np.arange(3))
      >>> square = mem.cache(np.square)
      >>> b = square(a)                                   # doctest: +ELLIPSIS
      ________________________________________________________________________________
      [Memory] Calling square...
      square(array([[0, 0, 1],
             [1, 1, 1],
             [4, 2, 1]]))
      ___________________________________________________________square - 0...s, 0.0min

      >>> c = square(a)
      >>> # The above call did not trigger an evaluation

2) **Embarrassingly parallel helper:** to make is easy to write readable
   parallel code and debug it quickly:

      >>> from joblib import Parallel, delayed
      >>> from math import sqrt
      >>> Parallel(n_jobs=1)(delayed(sqrt)(i**2) for i in range(10))
      [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]

3) **Logging/tracing:** The different functionalities will
   progressively acquire better logging mechanism to help track what
   has been ran, and capture I/O easily. In addition, Joblib will
   provide a few I/O primitives, to easily define define logging and
   display streams, and provide a way of compiling a report.
   We want to be able to quickly inspect what has been run.

..
    >>> import shutil ; shutil.rmtree('/tmp/joblib/')

"""

__version__ = '0.5.3'

from .memory import Memory
from .logger import PrintTime, Logger
from .hashing import hash
from .numpy_pickle import dump, load
from .parallel import Parallel, delayed, cpu_count
<|endoftext|>


<commit_before>"""
JobMon - Network Queue
======================

The JobMon supervisor has to manage two input sources - incoming connections
from clients, and changes in the state of its child processes. Information
about child processes is sent to the supervisor via a queue - however, queues
cannot be waited upon in conjunction with sockets.

This module 'converts' network traffic into queue elements. The queue elements
are pushed into the same queue which child events are pushed into, allowing the
supervisor to use a single input queue. In order for the supervisor to reply,
it pushes responses onto another queue - the thread which handles output queue
is also provided by this module.
"""
from collections import namedtuple
import logging
import os
import queue
import select
import socket
import threading

from jobmon import protocol, transport

# How long for threads to wait before checking their quit events
THREAD_TIMEOUT = 2

# A message received from a socket, or a message to send to a socket
SocketMessage = namedtuple('SocketMessage', ['message', 'socket'])

class NetworkCommandQueue:
    """
    Handles receiving and sending commands over the network.

    Incoming commands are read from the network, and are stored with the client
    socket into an incoming queue, which the supervisor should read from.
    Outgoing responses to commands should be pushed into the ougoing queue,
    where this thread will send the responses back out over the network.
    """
    def __init__(self, socket_path, net_input_queue):
        """
        Creates a new :class:`NetworkCommandQueue`.

        :param str socket_path: The path to the UNIX socket on which to accept\
        connections.
        :param queue.Queue net_input_queue: The queue to send network messages.

        Note that there is an additional queue, :attr:`net_output`, which is
        where network responses are read from. Note that both the input queue,
        and the output queue, deal only in terms of :class:`SocketMessage`
        objects.
        """
        self.sock_path = socket_path
        self.net_input = net_input_queue
        self.net_output = queue.Queue()
        self.quit_event = threading.Event()

        self.request_thread = None
        self.response_thread = None

        self.logger = logging.getLogger('supervisor.command-queue')

    def start(self):
        """
        Starts the network queue by launching the handler threads.
        """
        if self.request_thread or self.response_thread:
            raise ValueError('{} already started - cannot start again'.format(self))

        self.request_thread = threading.Thread(
            target=self.handle_network_requests)
        self.request_thread.daemon = True
        self.request_thread.start()

        self.response_thread = threading.Thread(
            target=self.handle_network_responses)
        self.response_thread.daemon = True
        self.response_thread.start()

    def stop(self):
        """
        Stops the network queue's handler threads and waits for them to exit.
        """
        if self.request_thread is None or self.response_thread is None:
            raise ValueError('{} not started'.format(self))

        self.quit_event.set()
        self.request_thread.join()
        self.response_thread.join()

    def handle_network_requests(self):
        """
        Handles requests from the network, which are then translated into queue
        messages.
        """
        server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server_socket.setblocking(False)
        server_socket.bind(self.sock_path)
        server_socket.listen(25)
        self.logger.info('Listening on %s for commands', server_socket)

        while True:
            # Since we need to know if the main thread wants to kill us, wait
            # for clients asynchronously, and handle the quit event if it is
            # sent to us
            connections, _, _ = select.select([server_socket], [], [],
                                              THREAD_TIMEOUT)

            if connections:
                client, _ = server_socket.accept()
                client.setblocking(False)

                # If the client drops, or is otherwise unable to get us the
                # data we want, then drop them to avoid getting hung up
                try:
                    self.logger.info('Reading message...')
                    message = protocol.recv_message(client)
                    self.logger.info('Received %s on %s', message, client)
                    self.net_input.put(SocketMessage(message, client))
                except OSError:
                    self.logger.info('Client disconnected')
                    client.close()

            if self.quit_event.isSet():
                break

        server_socket.close()
        os.remove(self.sock_path)

    def handle_network_responses(self):
        """
        Pushes responses from the supervisor back onto the network.
        """
        while True:
            try:
                request = self.net_output.get(timeout=THREAD_TIMEOUT)
                protocol.send_message(request.message, request.socket)
                self.logger.info('Sending %s to %s', request.message,
                                 request.socket)
                request.socket.close()
            except queue.Empty:
                pass

            if self.quit_event.isSet():
                break

class NetworkEventQueue:
    def __init__(self, socket_path):
        """
        Creates a new :class:`NetworkEventQueue`.

        :param str socket_path: The path to the UNIX socket on which to accept\
        connections.

        Note that there is an additional queue, :attr:`event_output`, which is
        where events should be sent to get them forwarded onto the network.

        If you want to wait for the server to set up its socket, then wait on
        :attr:`server_listening`.
        """
        self.sock_path = socket_path
        self.event_output = queue.Queue()
        self.quit_event = threading.Event()

        self.event_thread = None
        self.connection_thread = None
        self.connections = set()
        self.connections_modifier_lock = threading.Lock()
        self.server_listening = threading.Event()

        self.logger = logging.getLogger('supervisor.event_queue')

    def start(self):
        """
        Starts the network queue by launching the handler threads.
        """
        if self.event_thread or self.event_thread:
            raise ValueError('{} already started - cannot start again'.format(self))

        self.connection_thread = threading.Thread(target=self.handle_connections)
        self.connection_thread.daemon = True
        self.connection_thread.start()

        self.event_thread = threading.Thread(target=self.handle_events)
        self.event_thread.daemon = True
        self.event_thread.start()

    def stop(self):
        """
        Stops the handler thread and waits for it to exit.
        """
        if self.event_thread is None:
            raise ValueError('{} not started'.format(self))

        self.quit_event.set()
        self.event_thread.join()
        self.connection_thread.join()

        self.server_listening.clear()

    def handle_events(self):
        """
        Reads in events from a queue, and dispatches them to connected peers.
        """
        while True:
            try:
                event = self.event_output.get(timeout=THREAD_TIMEOUT)
                self.logger.info('Pushing out event to %d listeners', len(self.connections))

                with self.connections_modifier_lock:
                    for peer in self.connections:
                        try:
                            protocol.send_message(event, peer)
                            self.logger.info('- %s', peer)
                        except OSError:
                            # Since we're locking the connections set, dead
                            # connections can't be removed. Since the
                            # connection thread will remove them, we should
                            # ignore them
                            continue
            except queue.Empty:
                pass

            if self.quit_event.isSet():
                self.logger.info('Event thread exiting')
                break

    def handle_connections(self):
        """
        Handles incoming connections to the event socket, and drops them when
        they die.
        """
        server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server_socket.setblocking(False)
        server_socket.bind(self.sock_path)
        server_socket.listen(5)
        self.server_listening.set()
        self.logger.info('Dispatching events via %s', server_socket)

        while True:
            rlist, wlist, xlist = select.select([server_socket] + list(self.connections),
                                                [], [], THREAD_TIMEOUT)
            for reader in rlist:
                if reader is server_socket:
                    peer, _ = server_socket.accept()
                    self.connections.add(peer)
                    self.logger.info('Gained listener on %s', reader)
                else:
                    # If there's incoming data, it could be an EOF - if we read
                    # nothing, then the peer is gone
                    data = reader.recv(1)
                    if not data:
                        with self.connections_modifier_lock:
                            self.connections.remove(reader)
                            self.logger.info('Lost listener on %s', reader)

            if self.quit_event.isSet():
                self.logger.info('Connection thread closing')
                break

        server_socket.close()
        os.remove(self.sock_path)

        for peer in self.connections:
            peer.close()
        self.connections.clear()
<commit_msg>Made the event dispatcher withhold events.<commit_after>"""
JobMon - Network Queue
======================

The JobMon supervisor has to manage two input sources - incoming connections
from clients, and changes in the state of its child processes. Information
about child processes is sent to the supervisor via a queue - however, queues
cannot be waited upon in conjunction with sockets.

This module 'converts' network traffic into queue elements. The queue elements
are pushed into the same queue which child events are pushed into, allowing the
supervisor to use a single input queue. In order for the supervisor to reply,
it pushes responses onto another queue - the thread which handles output queue
is also provided by this module.
"""
from collections import deque, namedtuple
import logging
import os
import queue
import select
import socket
import threading

from jobmon import protocol, transport

# How long for threads to wait before checking their quit events
THREAD_TIMEOUT = 2

# A message received from a socket, or a message to send to a socket
SocketMessage = namedtuple('SocketMessage', ['message', 'socket'])

class NetworkCommandQueue:
    """
    Handles receiving and sending commands over the network.

    Incoming commands are read from the network, and are stored with the client
    socket into an incoming queue, which the supervisor should read from.
    Outgoing responses to commands should be pushed into the ougoing queue,
    where this thread will send the responses back out over the network.
    """
    def __init__(self, socket_path, net_input_queue):
        """
        Creates a new :class:`NetworkCommandQueue`.

        :param str socket_path: The path to the UNIX socket on which to accept\
        connections.
        :param queue.Queue net_input_queue: The queue to send network messages.

        Note that there is an additional queue, :attr:`net_output`, which is
        where network responses are read from. Note that both the input queue,
        and the output queue, deal only in terms of :class:`SocketMessage`
        objects.
        """
        self.sock_path = socket_path
        self.net_input = net_input_queue
        self.net_output = queue.Queue()
        self.quit_event = threading.Event()

        self.request_thread = None
        self.response_thread = None

        self.logger = logging.getLogger('supervisor.command-queue')

    def start(self):
        """
        Starts the network queue by launching the handler threads.
        """
        if self.request_thread or self.response_thread:
            raise ValueError('{} already started - cannot start again'.format(self))

        self.request_thread = threading.Thread(
            target=self.handle_network_requests)
        self.request_thread.daemon = True
        self.request_thread.start()

        self.response_thread = threading.Thread(
            target=self.handle_network_responses)
        self.response_thread.daemon = True
        self.response_thread.start()

    def stop(self):
        """
        Stops the network queue's handler threads and waits for them to exit.
        """
        if self.request_thread is None or self.response_thread is None:
            raise ValueError('{} not started'.format(self))

        self.quit_event.set()
        self.request_thread.join()
        self.response_thread.join()

    def handle_network_requests(self):
        """
        Handles requests from the network, which are then translated into queue
        messages.
        """
        server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server_socket.setblocking(False)
        server_socket.bind(self.sock_path)
        server_socket.listen(25)
        self.logger.info('Listening on %s for commands', server_socket)

        while True:
            # Since we need to know if the main thread wants to kill us, wait
            # for clients asynchronously, and handle the quit event if it is
            # sent to us
            connections, _, _ = select.select([server_socket], [], [],
                                              THREAD_TIMEOUT)

            if connections:
                client, _ = server_socket.accept()
                client.setblocking(False)

                # If the client drops, or is otherwise unable to get us the
                # data we want, then drop them to avoid getting hung up
                try:
                    self.logger.info('Reading message...')
                    message = protocol.recv_message(client)
                    self.logger.info('Received %s on %s', message, client)
                    self.net_input.put(SocketMessage(message, client))
                except OSError:
                    self.logger.info('Client disconnected')
                    client.close()

            if self.quit_event.isSet():
                break

        server_socket.close()
        os.remove(self.sock_path)

    def handle_network_responses(self):
        """
        Pushes responses from the supervisor back onto the network.
        """
        while True:
            try:
                request = self.net_output.get(timeout=THREAD_TIMEOUT)
                protocol.send_message(request.message, request.socket)
                self.logger.info('Sending %s to %s', request.message,
                                 request.socket)
                request.socket.close()
            except queue.Empty:
                pass

            if self.quit_event.isSet():
                break

class UndoableQueue(queue.Queue):
    """
    A queue which can have elements put back onto the front if they are not
    processed.
    """
    def __init__(self, maxsize=0, max_backlog=0):
        """
        Create a new :class:`UndoableQueue`.

        :param int maxsize: The maximum size of the queue, *not* including \
        the backlog.
        :param int max_backlog: The maximum size of the backlog, or 0 for an \
        unlimited size.
        """
        super().__init__(maxsize)

        # Ideally, this queue would be usable by multiple consumers/producers
        # (although this code doesn't need that capability), so we need to
        # protect the ungotten list.
        self.unget_lock = threading.Lock()

        # Where 'ungotten' elements are stored when they are put back onto the
        # head of the queue
        self.ungotten = deque()

        # The ungotten backlog can be of a finite size
        self.ungotten_size = max_backlog

    def unget(self, item):
        """
        Puts an element back onto the head of the queue.

        :param item: The object to put back onto the queue.
        """
        with self.unget_lock:
            if self.ungotten_size > 0 and len(self.ungotten) >= self.ungotten_size:
                # Rotate the first element out of the backlog, thus maintaining
                # the size when we add the next element
                self.ungotten.popleft()

            self.ungotten.append(item)

    def get(self, block=True, timeout=None):
        """
        Checks to see if any element can be removed from the queue, and returns
        that element if it does. See :meth:`queue.Queue.get` for the meaning of
        the other arguments.

        :return: The element on the head of the queue, and returns that \
        element if it does.
        """
        with self.unget_lock:
            if self.ungotten:
                return self.ungotten.popleft()
        return super().get(block, timeout)

class NetworkEventQueue:
    def __init__(self, socket_path):
        """
        Creates a new :class:`NetworkEventQueue`.

        :param str socket_path: The path to the UNIX socket on which to accept\
        connections.

        Note that there is an additional queue, :attr:`event_output`, which is
        where events should be sent to get them forwarded onto the network.

        If you want to wait for the server to set up its socket, then wait on
        :attr:`server_listening`.
        """
        self.sock_path = socket_path
        self.event_output = UndoableQueue()
        self.quit_event = threading.Event()

        self.event_thread = None
        self.connection_thread = None
        self.connections = set()
        self.connections_modifier_lock = threading.Lock()
        self.server_listening = threading.Event()

        self.logger = logging.getLogger('supervisor.event_queue')

    def start(self):
        """
        Starts the network queue by launching the handler threads.
        """
        if self.event_thread or self.event_thread:
            raise ValueError('{} already started - cannot start again'.format(self))

        self.connection_thread = threading.Thread(target=self.handle_connections)
        self.connection_thread.daemon = True
        self.connection_thread.start()

        self.event_thread = threading.Thread(target=self.handle_events)
        self.event_thread.daemon = True
        self.event_thread.start()

    def stop(self):
        """
        Stops the handler thread and waits for it to exit.
        """
        if self.event_thread is None:
            raise ValueError('{} not started'.format(self))

        self.quit_event.set()
        self.event_thread.join()
        self.connection_thread.join()

        self.server_listening.clear()

    def handle_events(self):
        """
        Reads in events from a queue, and dispatches them to connected peers.
        """
        while True:
            try:
                event = self.event_output.get(timeout=THREAD_TIMEOUT)

                with self.connections_modifier_lock:
                    if not self.connections:
                        self.logger.info('Keeping event in the queue, no listeners')
                        self.event_output.unget(event)
                    else:
                        self.logger.info('Pushing out event to %d listeners', len(self.connections))
                        for peer in self.connections:
                            try:
                                protocol.send_message(event, peer)
                                self.logger.info('- %s', peer)
                            except OSError:
                                # Since we're locking the connections set, dead
                                # connections can't be removed. Since the
                                # connection thread will remove them, we should
                                # ignore them
                                continue
            except queue.Empty:
                pass

            if self.quit_event.isSet():
                self.logger.info('Event thread exiting')
                break

    def handle_connections(self):
        """
        Handles incoming connections to the event socket, and drops them when
        they die.
        """
        server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server_socket.setblocking(False)
        server_socket.bind(self.sock_path)
        server_socket.listen(5)
        self.server_listening.set()
        self.logger.info('Dispatching events via %s', server_socket)

        while True:
            rlist, wlist, xlist = select.select([server_socket] + list(self.connections),
                                                [], [], THREAD_TIMEOUT)
            for reader in rlist:
                if reader is server_socket:
                    peer, _ = server_socket.accept()
                    self.connections.add(peer)
                    self.logger.info('Gained listener on %s', reader)
                else:
                    # If there's incoming data, it could be an EOF - if we read
                    # nothing, then the peer is gone
                    data = reader.recv(1)
                    if not data:
                        with self.connections_modifier_lock:
                            self.connections.remove(reader)
                            self.logger.info('Lost listener on %s', reader)

            if self.quit_event.isSet():
                self.logger.info('Connection thread closing')
                break

        server_socket.close()
        os.remove(self.sock_path)

        for peer in self.connections:
            peer.close()
        self.connections.clear()
<|endoftext|>


<commit_before>from typing import Any
from typing import Dict
from typing import Optional
from copy import deepcopy
from datetime import datetime
from datetime import timezone
import json

from blobstash.base.client import Client
from blobstash.base.iterator import BasePaginationIterator
from blobstash.docstore.attachment import add_attachment
from blobstash.docstore.attachment import fadd_attachment
from blobstash.docstore.attachment import get_attachment as get_attach
from blobstash.docstore.attachment import fget_attachment
from blobstash.docstore.attachment import Attachment
from blobstash.docstore.error import DocStoreError
from blobstash.docstore.query import Q  # noqa: unused-import
from blobstash.docstore.query import LogicalOperator
from blobstash.docstore.query import Not
from blobstash.docstore.query import LuaScript
from blobstash.docstore.query import LuaShortQuery
from blobstash.docstore.query import LuaStoredQuery
from blobstash.docstore.query import LuaShortQueryComplex
from blobstash.filetree import Node

import jsonpatch

# Keep a local cache of the docs to be able to generate a JSON Patch
_DOC_CACHE: Dict[str, "_Document"] = {}

class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Attachment):
            return obj.pointer
        else:
            return super().default(obj)

class MissingIDError(DocStoreError):
    """Error raised when the document does not contains a `ID`."""

class NotADocumentError(DocStoreError):
    """Error raises when the given document is not a `dict` instance."""

class _Document(dict):
    """Document is a dict subclass for document returned by the API, keep track of the ETag, and the document for the
    JSON Patch generation if needed."""

    def __setitem__(self, key, val) -> None:
        if key != "_id":
            self.checkpoint()
        dict.__setitem__(self, key, val)

    def checkpoint(self) -> None:
        """Force a checkpoint for the JSON Patch generation, can be use if the dict will be used to generate an object
        (the `__setitem__` will never get triggered this way)."""
        _id = self.get("_id")
        if _id is None:
            raise MissingIDError

        if _id not in _DOC_CACHE:
            doc = self.copy()
            del doc["_id"]
            _DOC_CACHE[_id] = deepcopy(doc)  # type: ignore

    def __repr__(self):
        return dict.__repr__(self)

class ID:
    """ID holds the document ID along with metadata."""

    def __init__(self, data) -> None:
        self._id = data.get("_id")
        self._created = data.get("_created")
        self._updated = data.get("_updated")
        self._version = data.get("_version")

    @classmethod
    def inject(cls, data):
        """Extracts ID infos from the document special keys and remove them, replacing
        `_id` with an instance of `ID`."""
        doc_id = cls(data)
        if doc_id._id:
            del data["_id"]
        if doc_id._created:
            del data["_created"]
        if doc_id._updated:
            del data["_updated"]
        if doc_id._version is not None:
            del data["_version"]
        data["_id"] = doc_id
        return doc_id

    def version(self) -> str:
        """Return the document version/ETag."""
        return self._version

    def id(self) -> str:
        """Return the document ID (hex-encoded)."""
        return self._id

    def created(self) -> Optional[datetime]:
        """Return a datetime representing the document creation date (i.e. the first version)."""
        return self._parse_dt(self._created)

    def updated(self) -> Optional[datetime]:
        """Return a datetime representing the document current version creation date."""
        return self._parse_dt(self._updated)

    def _parse_dt(self, dt_str: Optional[str]) -> Optional[datetime]:
        if dt_str is None:
            return None

        dt = datetime.strptime(dt_str, "%Y-%m-%dT%H:%M:%SZ")
        dt = dt.replace(tzinfo=timezone.utc)
        return dt.astimezone()

    def __hash__(self):
        return hash((self._version, self._id))

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False

        return (self._version, self._id) == (other.version(), other.id())

    def __ne__(self, other):
        if not isinstance(other, self.__class__):
            return False
        return not self.__eq__(other)

    def __repr__(self):
        return "blobstash.docstore.ID(_id={!r})".format(self._id)

    def __str__(self):
        return self._id

class DocVersionsIterator(BasePaginationIterator):
    def __init__(self, client, col_name, _id, params=None, limit=None, cursor=None):
        if isinstance(_id, ID):
            _id = _id.id()

        self._id = _id
        self.col_name = col_name
        super().__init__(
            client=client,
            path="/api/docstore/" + self.col_name + "/" + self._id + "/_versions",
            params=params,
            limit=limit,
            cursor=cursor,
        )

    def parse_data(self, resp):
        raw_docs = resp["data"]
        docs = []
        pointers = resp["pointers"]
        for doc in raw_docs:
            ID.inject(doc)
            _fill_pointers(doc, pointers)
            docs.append(_Document(doc))

        return docs

class DocsQueryIterator(BasePaginationIterator):
    def __init__(
        self,
        client,
        collection,
        query,
        script=None,
        stored_query=None,
        stored_query_args=None,
        as_of=None,
        params=None,
        limit=None,
        cursor=None,
        per_page=None,
    ):
        self.query = query
        self.script = None
        # TODO supprt stored query
        self.collection = collection
        self.as_of = as_of

        # XXX(tsileo): intelligent limit (i.e. limit=1000, but want to query them 100 by 100)
        # Handle raw Lua script
        if isinstance(query, LuaScript):
            script = query.script
            query = None
        elif isinstance(query, LuaStoredQuery):
            stored_query = query.name
            stored_query_args = json.dumps(query.args)

        # Handle default query operators
        elif isinstance(
            query, (LogicalOperator, Not, LuaShortQueryComplex, LuaShortQuery)
        ):
            query = str(query)
        else:
            if query:
                query = str(query)

        if isinstance(as_of, datetime):
            as_of = as_of.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")

        params = dict(
            query=query,
            script=script,
            stored_query_args=stored_query_args,
            stored_query=stored_query,
            as_of=as_of,
        )

        super().__init__(
            client=client,
            path="/api/docstore/" + self.collection.name,
            params=params,
            limit=limit,
            cursor=cursor,
            per_page=per_page,
        )

    def parse_data(self, resp):
        docs = []
        pointers = resp["pointers"]
        for raw_doc in resp["data"]:
            ID.inject(raw_doc)
            _fill_pointers(raw_doc, pointers)
            docs.append(_Document(raw_doc))
        return docs

def _fill_pointers(doc, pointers):
    """Replace the pointer by actual object representation."""
    for k, v in doc.items():
        if isinstance(v, str):
            if v.startswith("@filetree/ref:"):
                doc[k] = Attachment(v, Node.from_resp(pointers[v]))
        elif isinstance(v, dict):
            _fill_pointers(v, pointers)

class Collection:
    """Collection represents a collection (analog to a database)."""

    def __init__(self, client, name):
        self._client = client
        self.name = name

    def insert(self, doc):
        """Insert the given document."""
        if not isinstance(doc, dict):
            raise NotADocumentError

        if "_id" in doc and isinstance(doc["_id"], ID):
            return self.update(doc)

        # TODO(tsileo): bulk insert
        # TODO(tsileo): file attachment
        if isinstance(doc, list):
            for d in doc:
                self._insert(d)

        resp = self._client.request("POST", "/api/docstore/" + self.name, json=doc)
        doc_id = ID.inject(resp)

        doc["_id"] = doc_id
        rdoc = doc.copy()
        del rdoc["_id"]

        _DOC_CACHE[doc_id] = deepcopy(rdoc)

        return doc_id

    def update(self, doc):
        """Update the given document/list of documents."""
        _id = doc.get("_id")
        if _id is None:
            raise MissingIDError
        del doc["_id"]
        if _id in _DOC_CACHE:
            src = _DOC_CACHE[_id]
            pdoc = json.loads(json.dumps(doc, cls=JSONEncoder))
            p = jsonpatch.make_patch(src, pdoc)
            del _DOC_CACHE[_id]

            js = p.to_string()

            resp = self._client.request(
                "PATCH",
                "/api/docstore/" + self.name + "/" + _id.id(),
                headers={"If-Match": _id.version()},
                data=js,
            )
            doc_id = ID.inject(resp)
            doc["_id"] = doc_id
            rdoc = doc.copy()
            del rdoc["_id"]
            _DOC_CACHE[doc_id] = deepcopy(rdoc)
            return doc_id
            # FIXME(tsileo): catch status 412
        else:
            resp = self._client.request(
                "POST",
                "/api/docstore/" + self.name + "/" + _id.id(),
                headers={"If-Match": _id.version()},
                json=doc,
            )
            doc_id = ID.inject(resp)
            doc["_id"] = doc_id
            rdoc = doc.copy()
            del rdoc["_id"]
            _DOC_CACHE[doc_id] = deepcopy(rdoc)
            return doc_id

    def get_by_id(self, _id):
        """Fetch a document by its ID (string, a an `ID` instance)."""
        if isinstance(_id, ID):
            _id = _id.id()

        resp = self._client.request("GET", "/api/docstore/" + self.name + "/" + _id)
        doc = resp["data"]
        pointers = resp["pointers"]
        _id = ID.inject(doc)
        _fill_pointers(doc, pointers)

        return _Document(doc)

    def get_versions(self, _id):
        return DocVersionsIterator(self._client, self.name, _id)

        if isinstance(_id, ID):
            _id = _id.id()

        resp = self._client.request(
            "GET",
            "/api/docstore/" + self.name + "/" + _id + "/_versions",
            params=dict(limit=0),
        )
        raw_docs = resp["data"]
        docs = []
        pointers = resp["pointers"]
        for doc in raw_docs:
            _id = ID.inject(doc)
            _fill_pointers(doc, pointers)
            docs.append(_Document(doc))

        return docs

    def delete(self, doc_or_docs):
        """Delete the given document/list of document."""
        if isinstance(doc_or_docs, list):
            docs = [doc_or_docs]
        for doc in docs:
            if isinstance(doc, dict):
                try:
                    _id = doc["_id"]
                except KeyError:
                    raise MissingIDError
            elif isinstance(_id, ID):
                _id = _id.id()
            else:
                raise NotADocumentError

            self._client.request("DELETE", "/api/docstore/" + self.name + "/" + _id)

    def map_reduce(
        self, map_: LuaScript, reduce_: LuaScript, as_of: str = ""
    ) -> Dict[str, Dict[str, Any]]:
        payload = {"map": map_.script, "reduce": reduce_.script}
        resp = self._client.request(
            "POST", f"/api/docstore/{self.name}/_map_reduce?as_of={as_of}", json=payload
        )
        return resp["data"]

    def query(
        self,
        query=None,
        script=None,
        stored_query=None,
        stored_query_args=None,
        as_of=None,
        limit=None,
        cursor=None,
        per_page=None,
    ):
        """Query the collection and return an iterable cursor."""
        return DocsQueryIterator(
            self._client,
            self,
            query,
            script=script,
            stored_query=stored_query,
            stored_query_args=stored_query_args,
            as_of=as_of,
            limit=limit,
            cursor=cursor,
        )

    def get(self, query="", script=""):
        """Return the first document matching the query."""
        for doc in self.query(query, script=script, limit=1):
            return doc

        return None

    def __repr__(self):
        return "blobstash.docstore.Collection(name={!r})".format(self.name)

    def __str__(self):
        return self.__repr__()

class DocStoreClient:
    """BlobStash DocStore client."""

    def __init__(self, base_url: str = None, api_key: str = None) -> None:
        self._client = Client(
            base_url=base_url, api_key=api_key, json_encoder=JSONEncoder
        )

    def __getitem__(self, key):
        return self.collection(key)

    def __getattr__(self, name):
        return self.collection(name)

    def collection(self, name):
        """Returns a `Collection` instance for the given name."""
        return Collection(self._client, name)

    def collections(self):
        """Returns all the available collections."""
        collections = []
        resp = self._client.request("GET", "/api/docstore/")
        for col in resp["collections"]:
            collections.append(self.collection(col))
        return collections

    def fadd_attachment(self, name=None, fileobj=None, content_type=None):
        return fadd_attachment(self._client, name, fileobj, content_type)

    def add_attachment(self, path):
        """Upload the file/dir at path, and return the key to embed the file/dir as an attachment/filetree pointer.

        >>> doc['my_super_text_file'] = client.add_attachment('/path/to/my/text_file.txt')

        """
        return add_attachment(self._client, path)

    def fget_attachment(self, attachment):
        """Returns a fileobj (that needs to be closed) with the content off the attachment."""
        return fget_attachment(self._client, attachment)

    def get_attachment(self, attachment, path):
        return get_attach(self._client, attachment, path)

    def __repr__(self):
        return "blobstash.docstore.DocStoreClient(base_url={!r})".format(
            self._client.base_url
        )

    def __str__(self):
        return self.__repr__()
<commit_msg>docstore: improve pointers support<commit_after>from typing import Any
from typing import Dict
from typing import Optional
from copy import deepcopy
from datetime import datetime
from datetime import timezone
import json

from blobstash.base.client import Client
from blobstash.base.iterator import BasePaginationIterator
from blobstash.docstore.attachment import add_attachment
from blobstash.docstore.attachment import fadd_attachment
from blobstash.docstore.attachment import get_attachment as get_attach
from blobstash.docstore.attachment import fget_attachment
from blobstash.docstore.attachment import Attachment
from blobstash.docstore.error import DocStoreError
from blobstash.docstore.query import Q  # noqa: unused-import
from blobstash.docstore.query import LogicalOperator
from blobstash.docstore.query import Not
from blobstash.docstore.query import LuaScript
from blobstash.docstore.query import LuaShortQuery
from blobstash.docstore.query import LuaStoredQuery
from blobstash.docstore.query import LuaShortQueryComplex
from blobstash.filetree import Node

import jsonpatch

# Keep a local cache of the docs to be able to generate a JSON Patch
_DOC_CACHE: Dict[str, "_Document"] = {}

class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Attachment):
            return obj.pointer
        else:
            return super().default(obj)

class MissingIDError(DocStoreError):
    """Error raised when the document does not contains a `ID`."""

class NotADocumentError(DocStoreError):
    """Error raises when the given document is not a `dict` instance."""

class _Document(dict):
    """Document is a dict subclass for document returned by the API, keep track of the ETag, and the document for the
    JSON Patch generation if needed."""

    def __setitem__(self, key, val) -> None:
        if key != "_id":
            self.checkpoint()
        dict.__setitem__(self, key, val)

    def checkpoint(self) -> None:
        """Force a checkpoint for the JSON Patch generation, can be use if the dict will be used to generate an object
        (the `__setitem__` will never get triggered this way)."""
        _id = self.get("_id")
        if _id is None:
            raise MissingIDError

        if _id not in _DOC_CACHE:
            doc = self.copy()
            del doc["_id"]
            _DOC_CACHE[_id] = deepcopy(doc)  # type: ignore

    def __repr__(self):
        return dict.__repr__(self)

class ID:
    """ID holds the document ID along with metadata."""

    def __init__(self, data) -> None:
        self._id = data.get("_id")
        self._created = data.get("_created")
        self._updated = data.get("_updated")
        self._version = data.get("_version")

    @classmethod
    def inject(cls, data):
        """Extracts ID infos from the document special keys and remove them, replacing
        `_id` with an instance of `ID`."""
        doc_id = cls(data)
        if doc_id._id:
            del data["_id"]
        if doc_id._created:
            del data["_created"]
        if doc_id._updated:
            del data["_updated"]
        if doc_id._version is not None:
            del data["_version"]
        data["_id"] = doc_id
        return doc_id

    def version(self) -> str:
        """Return the document version/ETag."""
        return self._version

    def id(self) -> str:
        """Return the document ID (hex-encoded)."""
        return self._id

    def created(self) -> Optional[datetime]:
        """Return a datetime representing the document creation date (i.e. the first version)."""
        return self._parse_dt(self._created)

    def updated(self) -> Optional[datetime]:
        """Return a datetime representing the document current version creation date."""
        return self._parse_dt(self._updated)

    def _parse_dt(self, dt_str: Optional[str]) -> Optional[datetime]:
        if dt_str is None:
            return None

        dt = datetime.strptime(dt_str, "%Y-%m-%dT%H:%M:%SZ")
        dt = dt.replace(tzinfo=timezone.utc)
        return dt.astimezone()

    def __hash__(self):
        return hash((self._version, self._id))

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False

        return (self._version, self._id) == (other.version(), other.id())

    def __ne__(self, other):
        if not isinstance(other, self.__class__):
            return False
        return not self.__eq__(other)

    def __repr__(self):
        return "blobstash.docstore.ID(_id={!r})".format(self._id)

    def __str__(self):
        return self._id

class DocVersionsIterator(BasePaginationIterator):
    def __init__(self, client, col_name, _id, params=None, limit=None, cursor=None):
        if isinstance(_id, ID):
            _id = _id.id()

        self._id = _id
        self.col_name = col_name
        super().__init__(
            client=client,
            path="/api/docstore/" + self.col_name + "/" + self._id + "/_versions",
            params=params,
            limit=limit,
            cursor=cursor,
        )

    def parse_data(self, resp):
        raw_docs = resp["data"]
        docs = []
        pointers = resp["pointers"]
        for doc in raw_docs:
            ID.inject(doc)
            _fill_pointers(doc, pointers)
            docs.append(_Document(doc))

        return docs

class DocsQueryIterator(BasePaginationIterator):
    def __init__(
        self,
        client,
        collection,
        query,
        script=None,
        stored_query=None,
        stored_query_args=None,
        as_of=None,
        params=None,
        limit=None,
        cursor=None,
        per_page=None,
    ):
        self.query = query
        self.script = None
        # TODO supprt stored query
        self.collection = collection
        self.as_of = as_of

        # XXX(tsileo): intelligent limit (i.e. limit=1000, but want to query them 100 by 100)
        # Handle raw Lua script
        if isinstance(query, LuaScript):
            script = query.script
            query = None
        elif isinstance(query, LuaStoredQuery):
            stored_query = query.name
            stored_query_args = json.dumps(query.args)

        # Handle default query operators
        elif isinstance(
            query, (LogicalOperator, Not, LuaShortQueryComplex, LuaShortQuery)
        ):
            query = str(query)
        else:
            if query:
                query = str(query)

        if isinstance(as_of, datetime):
            as_of = as_of.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")

        params = dict(
            query=query,
            script=script,
            stored_query_args=stored_query_args,
            stored_query=stored_query,
            as_of=as_of,
        )

        super().__init__(
            client=client,
            path="/api/docstore/" + self.collection.name,
            params=params,
            limit=limit,
            cursor=cursor,
            per_page=per_page,
        )

    def parse_data(self, resp):
        docs = []
        pointers = resp["pointers"]
        for raw_doc in resp["data"]:
            ID.inject(raw_doc)
            _fill_pointers(raw_doc, pointers)
            docs.append(_Document(raw_doc))
        return docs

def _fill_pointers(doc, pointers):
    """Replace the pointer by actual object representation."""
    if not isinstance(doc, dict):
        return

    for k, v in doc.items():
        if isinstance(v, str):
            if v.startswith("@filetree/ref:"):
                doc[k] = Attachment(v, Node.from_resp(pointers[v]))
        elif isinstance(v, dict):
            _fill_pointers(v, pointers)
        elif isinstance(v, list):
            doc[k] = [
                Attachment(item, Node.from_resp(pointers[item]))
                if isinstance(item, str) and item.startswith("@filetree/ref:")
                else item
                for item in v
            ]
            for item in doc[k]:
                _fill_pointers(item, pointers)

class Collection:
    """Collection represents a collection (analog to a database)."""

    def __init__(self, client, name):
        self._client = client
        self.name = name

    def insert(self, doc):
        """Insert the given document."""
        if not isinstance(doc, dict):
            raise NotADocumentError

        if "_id" in doc and isinstance(doc["_id"], ID):
            return self.update(doc)

        # TODO(tsileo): bulk insert
        # TODO(tsileo): file attachment
        if isinstance(doc, list):
            for d in doc:
                self._insert(d)

        resp = self._client.request("POST", "/api/docstore/" + self.name, json=doc)
        doc_id = ID.inject(resp)

        doc["_id"] = doc_id
        rdoc = doc.copy()
        del rdoc["_id"]

        _DOC_CACHE[doc_id] = deepcopy(rdoc)

        return doc_id

    def update(self, doc):
        """Update the given document/list of documents."""
        _id = doc.get("_id")
        if _id is None:
            raise MissingIDError
        del doc["_id"]
        if _id in _DOC_CACHE:
            src = _DOC_CACHE[_id]
            pdoc = json.loads(json.dumps(doc, cls=JSONEncoder))
            p = jsonpatch.make_patch(src, pdoc)
            del _DOC_CACHE[_id]

            js = p.to_string()

            resp = self._client.request(
                "PATCH",
                "/api/docstore/" + self.name + "/" + _id.id(),
                headers={"If-Match": _id.version()},
                data=js,
            )
            doc_id = ID.inject(resp)
            doc["_id"] = doc_id
            rdoc = doc.copy()
            del rdoc["_id"]
            _DOC_CACHE[doc_id] = deepcopy(rdoc)
            return doc_id
            # FIXME(tsileo): catch status 412
        else:
            resp = self._client.request(
                "POST",
                "/api/docstore/" + self.name + "/" + _id.id(),
                headers={"If-Match": _id.version()},
                json=doc,
            )
            doc_id = ID.inject(resp)
            doc["_id"] = doc_id
            rdoc = doc.copy()
            del rdoc["_id"]
            _DOC_CACHE[doc_id] = deepcopy(rdoc)
            return doc_id

    def get_by_id(self, _id):
        """Fetch a document by its ID (string, a an `ID` instance)."""
        if isinstance(_id, ID):
            _id = _id.id()

        resp = self._client.request("GET", "/api/docstore/" + self.name + "/" + _id)
        doc = resp["data"]
        pointers = resp["pointers"]
        _id = ID.inject(doc)
        _fill_pointers(doc, pointers)

        return _Document(doc)

    def get_versions(self, _id):
        return DocVersionsIterator(self._client, self.name, _id)

        if isinstance(_id, ID):
            _id = _id.id()

        resp = self._client.request(
            "GET",
            "/api/docstore/" + self.name + "/" + _id + "/_versions",
            params=dict(limit=0),
        )
        raw_docs = resp["data"]
        docs = []
        pointers = resp["pointers"]
        for doc in raw_docs:
            _id = ID.inject(doc)
            _fill_pointers(doc, pointers)
            docs.append(_Document(doc))

        return docs

    def delete(self, doc_or_docs):
        """Delete the given document/list of document."""
        if isinstance(doc_or_docs, list):
            docs = [doc_or_docs]
        for doc in docs:
            if isinstance(doc, dict):
                try:
                    _id = doc["_id"]
                except KeyError:
                    raise MissingIDError
            elif isinstance(_id, ID):
                _id = _id.id()
            else:
                raise NotADocumentError

            self._client.request("DELETE", "/api/docstore/" + self.name + "/" + _id)

    def map_reduce(
        self, map_: LuaScript, reduce_: LuaScript, as_of: str = ""
    ) -> Dict[str, Dict[str, Any]]:
        payload = {"map": map_.script, "reduce": reduce_.script}
        resp = self._client.request(
            "POST", f"/api/docstore/{self.name}/_map_reduce?as_of={as_of}", json=payload
        )
        return resp["data"]

    def query(
        self,
        query=None,
        script=None,
        stored_query=None,
        stored_query_args=None,
        as_of=None,
        limit=None,
        cursor=None,
        per_page=None,
    ):
        """Query the collection and return an iterable cursor."""
        return DocsQueryIterator(
            self._client,
            self,
            query,
            script=script,
            stored_query=stored_query,
            stored_query_args=stored_query_args,
            as_of=as_of,
            limit=limit,
            cursor=cursor,
        )

    def get(self, query="", script=""):
        """Return the first document matching the query."""
        for doc in self.query(query, script=script, limit=1):
            return doc

        return None

    def __repr__(self):
        return "blobstash.docstore.Collection(name={!r})".format(self.name)

    def __str__(self):
        return self.__repr__()

class DocStoreClient:
    """BlobStash DocStore client."""

    def __init__(self, base_url: str = None, api_key: str = None) -> None:
        self._client = Client(
            base_url=base_url, api_key=api_key, json_encoder=JSONEncoder
        )

    def __getitem__(self, key):
        return self.collection(key)

    def __getattr__(self, name):
        return self.collection(name)

    def collection(self, name):
        """Returns a `Collection` instance for the given name."""
        return Collection(self._client, name)

    def collections(self):
        """Returns all the available collections."""
        collections = []
        resp = self._client.request("GET", "/api/docstore/")
        for col in resp["collections"]:
            collections.append(self.collection(col))
        return collections

    def fadd_attachment(self, name=None, fileobj=None, content_type=None):
        return fadd_attachment(self._client, name, fileobj, content_type)

    def add_attachment(self, path):
        """Upload the file/dir at path, and return the key to embed the file/dir as an attachment/filetree pointer.

        >>> doc['my_super_text_file'] = client.add_attachment('/path/to/my/text_file.txt')

        """
        return add_attachment(self._client, path)

    def fget_attachment(self, attachment):
        """Returns a fileobj (that needs to be closed) with the content off the attachment."""
        return fget_attachment(self._client, attachment)

    def get_attachment(self, attachment, path):
        return get_attach(self._client, attachment, path)

    def __repr__(self):
        return "blobstash.docstore.DocStoreClient(base_url={!r})".format(
            self._client.base_url
        )

    def __str__(self):
        return self.__repr__()
<|endoftext|>