<commit_before># coding: utf-8
#
# Copyright 2011 Yesudeep Mangalapilly <yesudeep@gmail.com>
# Copyright 2012 Google, Inc & contributors.
#
# 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.

"""
:module: watchdog.events
:synopsis: File system events and event handlers.
:author: yesudeep@google.com (Yesudeep Mangalapilly)
:author: contact@tiger-222.fr (Mickaël Schoentgen)

Event Classes
-------------
.. autoclass:: FileSystemEvent
   :members:
   :show-inheritance:
   :inherited-members:

.. autoclass:: FileSystemMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileModifiedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirModifiedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileCreatedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileClosedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirCreatedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileDeletedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirDeletedEvent
   :members:
   :show-inheritance:

Event Handler Classes
---------------------
.. autoclass:: FileSystemEventHandler
   :members:
   :show-inheritance:

.. autoclass:: PatternMatchingEventHandler
   :members:
   :show-inheritance:

.. autoclass:: RegexMatchingEventHandler
   :members:
   :show-inheritance:

.. autoclass:: LoggingEventHandler
   :members:
   :show-inheritance:

"""

import os.path
import logging
import re
from watchdog.utils.patterns import match_any_paths

EVENT_TYPE_MOVED = 'moved'
EVENT_TYPE_DELETED = 'deleted'
EVENT_TYPE_CREATED = 'created'
EVENT_TYPE_MODIFIED = 'modified'
EVENT_TYPE_CLOSED = 'closed'

class FileSystemEvent:
    """
    Immutable type that represents a file system event that is triggered
    when a change occurs on the monitored file system.

    All FileSystemEvent objects are required to be immutable and hence
    can be used as keys in dictionaries or be added to sets.
    """

    event_type = None
    """The type of the event as a string."""

    is_directory = False
    """True if event was emitted for a directory; False otherwise."""

    is_synthetic = False
    """
    True if event was synthesized; False otherwise.

    These are events that weren't actually broadcast by the OS, but
    are presumed to have happened based on other, actual events.
    """

    def __init__(self, src_path):
        self._src_path = src_path

    @property
    def src_path(self):
        """Source path of the file system object that triggered this event."""
        return self._src_path

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

    def __repr__(self):
        return ("<%(class_name)s: event_type=%(event_type)s, "
                "src_path=%(src_path)r, "
                "is_directory=%(is_directory)s>"
                ) % (dict(
                     class_name=self.__class__.__name__,
                     event_type=self.event_type,
                     src_path=self.src_path,
                     is_directory=self.is_directory))

    # Used for comparison of events.
    @property
    def key(self):
        return (self.event_type, self.src_path, self.is_directory)

    def __eq__(self, event):
        return self.key == event.key

    def __ne__(self, event):
        return self.key != event.key

    def __hash__(self):
        return hash(self.key)

class FileSystemMovedEvent(FileSystemEvent):
    """
    File system event representing any kind of file system movement.
    """

    event_type = EVENT_TYPE_MOVED

    def __init__(self, src_path, dest_path):
        super().__init__(src_path)
        self._dest_path = dest_path

    @property
    def dest_path(self):
        """The destination path of the move event."""
        return self._dest_path

    # Used for hashing this as an immutable object.
    @property
    def key(self):
        return (self.event_type, self.src_path, self.dest_path, self.is_directory)

    def __repr__(self):
        return ("<%(class_name)s: src_path=%(src_path)r, "
                "dest_path=%(dest_path)r, "
                "is_directory=%(is_directory)s>"
                ) % (dict(class_name=self.__class__.__name__,
                          src_path=self.src_path,
                          dest_path=self.dest_path,
                          is_directory=self.is_directory))

# File events.

class FileDeletedEvent(FileSystemEvent):
    """File system event representing file deletion on the file system."""

    event_type = EVENT_TYPE_DELETED

class FileModifiedEvent(FileSystemEvent):
    """File system event representing file modification on the file system."""

    event_type = EVENT_TYPE_MODIFIED

class FileCreatedEvent(FileSystemEvent):
    """File system event representing file creation on the file system."""

    event_type = EVENT_TYPE_CREATED

class FileMovedEvent(FileSystemMovedEvent):
    """File system event representing file movement on the file system."""

class FileClosedEvent(FileSystemEvent):
    """File system event representing file close on the file system."""

    event_type = EVENT_TYPE_CLOSED

# Directory events.

class DirDeletedEvent(FileSystemEvent):
    """File system event representing directory deletion on the file system."""

    event_type = EVENT_TYPE_DELETED
    is_directory = True

class DirModifiedEvent(FileSystemEvent):
    """
    File system event representing directory modification on the file system.
    """

    event_type = EVENT_TYPE_MODIFIED
    is_directory = True

class DirCreatedEvent(FileSystemEvent):
    """File system event representing directory creation on the file system."""

    event_type = EVENT_TYPE_CREATED
    is_directory = True

class DirMovedEvent(FileSystemMovedEvent):
    """File system event representing directory movement on the file system."""

    is_directory = True

class FileSystemEventHandler:
    """
    Base file system event handler that you can override methods from.
    """

    def __init__(self):
        # Allow overriding or adding custom event handlers to event dispatch map.
        self.event_dispatch_map = {
            EVENT_TYPE_CREATED: self.on_created,
            EVENT_TYPE_DELETED: self.on_deleted,
            EVENT_TYPE_MODIFIED: self.on_modified,
            EVENT_TYPE_MOVED: self.on_moved,
            EVENT_TYPE_CLOSED: self.on_closed,
        }

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        self.on_any_event(event)
        self.event_dispatch_map[event.event_type](event)

    def on_any_event(self, event):
        """Catch-all event handler.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """

    def on_moved(self, event):
        """Called when a file or a directory is moved or renamed.

        :param event:
            Event representing file/directory movement.
        :type event:
            :class:`DirMovedEvent` or :class:`FileMovedEvent`
        """

    def on_created(self, event):
        """Called when a file or directory is created.

        :param event:
            Event representing file/directory creation.
        :type event:
            :class:`DirCreatedEvent` or :class:`FileCreatedEvent`
        """

    def on_deleted(self, event):
        """Called when a file or directory is deleted.

        :param event:
            Event representing file/directory deletion.
        :type event:
            :class:`DirDeletedEvent` or :class:`FileDeletedEvent`
        """

    def on_modified(self, event):
        """Called when a file or directory is modified.

        :param event:
            Event representing file/directory modification.
        :type event:
            :class:`DirModifiedEvent` or :class:`FileModifiedEvent`
        """

    def on_closed(self, event):
        """Called when a file opened for writing is closed.

        :param event:
            Event representing file closing.
        :type event:
            :class:`FileClosedEvent`
        """

class PatternMatchingEventHandler(FileSystemEventHandler):
    """
    Matches given patterns with file paths associated with occurring events.
    """

    def __init__(self, patterns=None, ignore_patterns=None,
                 ignore_directories=False, case_sensitive=False):
        super().__init__()

        self._patterns = patterns
        self._ignore_patterns = ignore_patterns
        self._ignore_directories = ignore_directories
        self._case_sensitive = case_sensitive

    @property
    def patterns(self):
        """
        (Read-only)
        Patterns to allow matching event paths.
        """
        return self._patterns

    @property
    def ignore_patterns(self):
        """
        (Read-only)
        Patterns to ignore matching event paths.
        """
        return self._ignore_patterns

    @property
    def ignore_directories(self):
        """
        (Read-only)
        ``True`` if directories should be ignored; ``False`` otherwise.
        """
        return self._ignore_directories

    @property
    def case_sensitive(self):
        """
        (Read-only)
        ``True`` if path names should be matched sensitive to case; ``False``
        otherwise.
        """
        return self._case_sensitive

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        if self.ignore_directories and event.is_directory:
            return

        paths = []
        if hasattr(event, 'dest_path'):
            paths.append(os.fsdecode(event.dest_path))
        if event.src_path:
            paths.append(os.fsdecode(event.src_path))

        if match_any_paths(paths,
                           included_patterns=self.patterns,
                           excluded_patterns=self.ignore_patterns,
                           case_sensitive=self.case_sensitive):
            super().dispatch(event)

class RegexMatchingEventHandler(FileSystemEventHandler):
    """
    Matches given regexes with file paths associated with occurring events.
    """

    def __init__(self, regexes=None, ignore_regexes=None,
                 ignore_directories=False, case_sensitive=False):
        super().__init__()

        if regexes is None:
            regexes = [r".*"]
        elif isinstance(regexes, str):
            regexes = [regexes]
        if ignore_regexes is None:
            ignore_regexes = []
        if case_sensitive:
            self._regexes = [re.compile(r) for r in regexes]
            self._ignore_regexes = [re.compile(r) for r in ignore_regexes]
        else:
            self._regexes = [re.compile(r, re.I) for r in regexes]
            self._ignore_regexes = [re.compile(r, re.I) for r in ignore_regexes]
        self._ignore_directories = ignore_directories
        self._case_sensitive = case_sensitive

    @property
    def regexes(self):
        """
        (Read-only)
        Regexes to allow matching event paths.
        """
        return self._regexes

    @property
    def ignore_regexes(self):
        """
        (Read-only)
        Regexes to ignore matching event paths.
        """
        return self._ignore_regexes

    @property
    def ignore_directories(self):
        """
        (Read-only)
        ``True`` if directories should be ignored; ``False`` otherwise.
        """
        return self._ignore_directories

    @property
    def case_sensitive(self):
        """
        (Read-only)
        ``True`` if path names should be matched sensitive to case; ``False``
        otherwise.
        """
        return self._case_sensitive

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        if self.ignore_directories and event.is_directory:
            return

        paths = []
        if hasattr(event, 'dest_path'):
            paths.append(os.fsdecode(event.dest_path))
        if event.src_path:
            paths.append(os.fsdecode(event.src_path))

        if any(r.match(p) for r in self.ignore_regexes for p in paths):
            return

        if any(r.match(p) for r in self.regexes for p in paths):
            super().dispatch(event)

class LoggingEventHandler(FileSystemEventHandler):
    """Logs all the events captured."""

    def __init__(self, logger=None):
        super().__init__()

        self.logger = logger or logging.root

    def on_moved(self, event):
        super().on_moved(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Moved %s: from %s to %s", what, event.src_path,
                         event.dest_path)

    def on_created(self, event):
        super().on_created(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Created %s: %s", what, event.src_path)

    def on_deleted(self, event):
        super().on_deleted(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Deleted %s: %s", what, event.src_path)

    def on_modified(self, event):
        super().on_modified(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Modified %s: %s", what, event.src_path)

def generate_sub_moved_events(src_dir_path, dest_dir_path):
    """Generates an event list of :class:`DirMovedEvent` and
    :class:`FileMovedEvent` objects for all the files and directories within
    the given moved directory that were moved along with the directory.

    :param src_dir_path:
        The source path of the moved directory.
    :param dest_dir_path:
        The destination path of the moved directory.
    :returns:
        An iterable of file system events of type :class:`DirMovedEvent` and
        :class:`FileMovedEvent`.
    """
    for root, directories, filenames in os.walk(dest_dir_path):
        for directory in directories:
            full_path = os.path.join(root, directory)
            renamed_path = full_path.replace(dest_dir_path, src_dir_path) if src_dir_path else None
            event = DirMovedEvent(renamed_path, full_path)
            event.is_synthetic = True
            yield event
        for filename in filenames:
            full_path = os.path.join(root, filename)
            renamed_path = full_path.replace(dest_dir_path, src_dir_path) if src_dir_path else None
            event = FileMovedEvent(renamed_path, full_path)
            event.is_synthetic = True
            yield event

def generate_sub_created_events(src_dir_path):
    """Generates an event list of :class:`DirCreatedEvent` and
    :class:`FileCreatedEvent` objects for all the files and directories within
    the given moved directory that were moved along with the directory.

    :param src_dir_path:
        The source path of the created directory.
    :returns:
        An iterable of file system events of type :class:`DirCreatedEvent` and
        :class:`FileCreatedEvent`.
    """
    for root, directories, filenames in os.walk(src_dir_path):
        for directory in directories:
            event = DirCreatedEvent(os.path.join(root, directory))
            event.is_synthetic = True
            yield event
        for filename in filenames:
            event = FileCreatedEvent(os.path.join(root, filename))
            event.is_synthetic = True
            yield event
<commit_msg>Revert "Allow overriding or adding custom event handlers to event dispatch map (#814)"<commit_after># coding: utf-8
#
# Copyright 2011 Yesudeep Mangalapilly <yesudeep@gmail.com>
# Copyright 2012 Google, Inc & contributors.
#
# 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.

"""
:module: watchdog.events
:synopsis: File system events and event handlers.
:author: yesudeep@google.com (Yesudeep Mangalapilly)
:author: contact@tiger-222.fr (Mickaël Schoentgen)

Event Classes
-------------
.. autoclass:: FileSystemEvent
   :members:
   :show-inheritance:
   :inherited-members:

.. autoclass:: FileSystemMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirMovedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileModifiedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirModifiedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileCreatedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileClosedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirCreatedEvent
   :members:
   :show-inheritance:

.. autoclass:: FileDeletedEvent
   :members:
   :show-inheritance:

.. autoclass:: DirDeletedEvent
   :members:
   :show-inheritance:

Event Handler Classes
---------------------
.. autoclass:: FileSystemEventHandler
   :members:
   :show-inheritance:

.. autoclass:: PatternMatchingEventHandler
   :members:
   :show-inheritance:

.. autoclass:: RegexMatchingEventHandler
   :members:
   :show-inheritance:

.. autoclass:: LoggingEventHandler
   :members:
   :show-inheritance:

"""

import os.path
import logging
import re
from watchdog.utils.patterns import match_any_paths

EVENT_TYPE_MOVED = 'moved'
EVENT_TYPE_DELETED = 'deleted'
EVENT_TYPE_CREATED = 'created'
EVENT_TYPE_MODIFIED = 'modified'
EVENT_TYPE_CLOSED = 'closed'

class FileSystemEvent:
    """
    Immutable type that represents a file system event that is triggered
    when a change occurs on the monitored file system.

    All FileSystemEvent objects are required to be immutable and hence
    can be used as keys in dictionaries or be added to sets.
    """

    event_type = None
    """The type of the event as a string."""

    is_directory = False
    """True if event was emitted for a directory; False otherwise."""

    is_synthetic = False
    """
    True if event was synthesized; False otherwise.

    These are events that weren't actually broadcast by the OS, but
    are presumed to have happened based on other, actual events.
    """

    def __init__(self, src_path):
        self._src_path = src_path

    @property
    def src_path(self):
        """Source path of the file system object that triggered this event."""
        return self._src_path

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

    def __repr__(self):
        return ("<%(class_name)s: event_type=%(event_type)s, "
                "src_path=%(src_path)r, "
                "is_directory=%(is_directory)s>"
                ) % (dict(
                     class_name=self.__class__.__name__,
                     event_type=self.event_type,
                     src_path=self.src_path,
                     is_directory=self.is_directory))

    # Used for comparison of events.
    @property
    def key(self):
        return (self.event_type, self.src_path, self.is_directory)

    def __eq__(self, event):
        return self.key == event.key

    def __ne__(self, event):
        return self.key != event.key

    def __hash__(self):
        return hash(self.key)

class FileSystemMovedEvent(FileSystemEvent):
    """
    File system event representing any kind of file system movement.
    """

    event_type = EVENT_TYPE_MOVED

    def __init__(self, src_path, dest_path):
        super().__init__(src_path)
        self._dest_path = dest_path

    @property
    def dest_path(self):
        """The destination path of the move event."""
        return self._dest_path

    # Used for hashing this as an immutable object.
    @property
    def key(self):
        return (self.event_type, self.src_path, self.dest_path, self.is_directory)

    def __repr__(self):
        return ("<%(class_name)s: src_path=%(src_path)r, "
                "dest_path=%(dest_path)r, "
                "is_directory=%(is_directory)s>"
                ) % (dict(class_name=self.__class__.__name__,
                          src_path=self.src_path,
                          dest_path=self.dest_path,
                          is_directory=self.is_directory))

# File events.

class FileDeletedEvent(FileSystemEvent):
    """File system event representing file deletion on the file system."""

    event_type = EVENT_TYPE_DELETED

class FileModifiedEvent(FileSystemEvent):
    """File system event representing file modification on the file system."""

    event_type = EVENT_TYPE_MODIFIED

class FileCreatedEvent(FileSystemEvent):
    """File system event representing file creation on the file system."""

    event_type = EVENT_TYPE_CREATED

class FileMovedEvent(FileSystemMovedEvent):
    """File system event representing file movement on the file system."""

class FileClosedEvent(FileSystemEvent):
    """File system event representing file close on the file system."""

    event_type = EVENT_TYPE_CLOSED

# Directory events.

class DirDeletedEvent(FileSystemEvent):
    """File system event representing directory deletion on the file system."""

    event_type = EVENT_TYPE_DELETED
    is_directory = True

class DirModifiedEvent(FileSystemEvent):
    """
    File system event representing directory modification on the file system.
    """

    event_type = EVENT_TYPE_MODIFIED
    is_directory = True

class DirCreatedEvent(FileSystemEvent):
    """File system event representing directory creation on the file system."""

    event_type = EVENT_TYPE_CREATED
    is_directory = True

class DirMovedEvent(FileSystemMovedEvent):
    """File system event representing directory movement on the file system."""

    is_directory = True

class FileSystemEventHandler:
    """
    Base file system event handler that you can override methods from.
    """

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        self.on_any_event(event)
        {
            EVENT_TYPE_CREATED: self.on_created,
            EVENT_TYPE_DELETED: self.on_deleted,
            EVENT_TYPE_MODIFIED: self.on_modified,
            EVENT_TYPE_MOVED: self.on_moved,
            EVENT_TYPE_CLOSED: self.on_closed,
        }[event.event_type](event)

    def on_any_event(self, event):
        """Catch-all event handler.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """

    def on_moved(self, event):
        """Called when a file or a directory is moved or renamed.

        :param event:
            Event representing file/directory movement.
        :type event:
            :class:`DirMovedEvent` or :class:`FileMovedEvent`
        """

    def on_created(self, event):
        """Called when a file or directory is created.

        :param event:
            Event representing file/directory creation.
        :type event:
            :class:`DirCreatedEvent` or :class:`FileCreatedEvent`
        """

    def on_deleted(self, event):
        """Called when a file or directory is deleted.

        :param event:
            Event representing file/directory deletion.
        :type event:
            :class:`DirDeletedEvent` or :class:`FileDeletedEvent`
        """

    def on_modified(self, event):
        """Called when a file or directory is modified.

        :param event:
            Event representing file/directory modification.
        :type event:
            :class:`DirModifiedEvent` or :class:`FileModifiedEvent`
        """

    def on_closed(self, event):
        """Called when a file opened for writing is closed.

        :param event:
            Event representing file closing.
        :type event:
            :class:`FileClosedEvent`
        """

class PatternMatchingEventHandler(FileSystemEventHandler):
    """
    Matches given patterns with file paths associated with occurring events.
    """

    def __init__(self, patterns=None, ignore_patterns=None,
                 ignore_directories=False, case_sensitive=False):
        super().__init__()

        self._patterns = patterns
        self._ignore_patterns = ignore_patterns
        self._ignore_directories = ignore_directories
        self._case_sensitive = case_sensitive

    @property
    def patterns(self):
        """
        (Read-only)
        Patterns to allow matching event paths.
        """
        return self._patterns

    @property
    def ignore_patterns(self):
        """
        (Read-only)
        Patterns to ignore matching event paths.
        """
        return self._ignore_patterns

    @property
    def ignore_directories(self):
        """
        (Read-only)
        ``True`` if directories should be ignored; ``False`` otherwise.
        """
        return self._ignore_directories

    @property
    def case_sensitive(self):
        """
        (Read-only)
        ``True`` if path names should be matched sensitive to case; ``False``
        otherwise.
        """
        return self._case_sensitive

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        if self.ignore_directories and event.is_directory:
            return

        paths = []
        if hasattr(event, 'dest_path'):
            paths.append(os.fsdecode(event.dest_path))
        if event.src_path:
            paths.append(os.fsdecode(event.src_path))

        if match_any_paths(paths,
                           included_patterns=self.patterns,
                           excluded_patterns=self.ignore_patterns,
                           case_sensitive=self.case_sensitive):
            super().dispatch(event)

class RegexMatchingEventHandler(FileSystemEventHandler):
    """
    Matches given regexes with file paths associated with occurring events.
    """

    def __init__(self, regexes=None, ignore_regexes=None,
                 ignore_directories=False, case_sensitive=False):
        super().__init__()

        if regexes is None:
            regexes = [r".*"]
        elif isinstance(regexes, str):
            regexes = [regexes]
        if ignore_regexes is None:
            ignore_regexes = []
        if case_sensitive:
            self._regexes = [re.compile(r) for r in regexes]
            self._ignore_regexes = [re.compile(r) for r in ignore_regexes]
        else:
            self._regexes = [re.compile(r, re.I) for r in regexes]
            self._ignore_regexes = [re.compile(r, re.I) for r in ignore_regexes]
        self._ignore_directories = ignore_directories
        self._case_sensitive = case_sensitive

    @property
    def regexes(self):
        """
        (Read-only)
        Regexes to allow matching event paths.
        """
        return self._regexes

    @property
    def ignore_regexes(self):
        """
        (Read-only)
        Regexes to ignore matching event paths.
        """
        return self._ignore_regexes

    @property
    def ignore_directories(self):
        """
        (Read-only)
        ``True`` if directories should be ignored; ``False`` otherwise.
        """
        return self._ignore_directories

    @property
    def case_sensitive(self):
        """
        (Read-only)
        ``True`` if path names should be matched sensitive to case; ``False``
        otherwise.
        """
        return self._case_sensitive

    def dispatch(self, event):
        """Dispatches events to the appropriate methods.

        :param event:
            The event object representing the file system event.
        :type event:
            :class:`FileSystemEvent`
        """
        if self.ignore_directories and event.is_directory:
            return

        paths = []
        if hasattr(event, 'dest_path'):
            paths.append(os.fsdecode(event.dest_path))
        if event.src_path:
            paths.append(os.fsdecode(event.src_path))

        if any(r.match(p) for r in self.ignore_regexes for p in paths):
            return

        if any(r.match(p) for r in self.regexes for p in paths):
            super().dispatch(event)

class LoggingEventHandler(FileSystemEventHandler):
    """Logs all the events captured."""

    def __init__(self, logger=None):
        super().__init__()

        self.logger = logger or logging.root

    def on_moved(self, event):
        super().on_moved(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Moved %s: from %s to %s", what, event.src_path,
                         event.dest_path)

    def on_created(self, event):
        super().on_created(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Created %s: %s", what, event.src_path)

    def on_deleted(self, event):
        super().on_deleted(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Deleted %s: %s", what, event.src_path)

    def on_modified(self, event):
        super().on_modified(event)

        what = 'directory' if event.is_directory else 'file'
        self.logger.info("Modified %s: %s", what, event.src_path)

def generate_sub_moved_events(src_dir_path, dest_dir_path):
    """Generates an event list of :class:`DirMovedEvent` and
    :class:`FileMovedEvent` objects for all the files and directories within
    the given moved directory that were moved along with the directory.

    :param src_dir_path:
        The source path of the moved directory.
    :param dest_dir_path:
        The destination path of the moved directory.
    :returns:
        An iterable of file system events of type :class:`DirMovedEvent` and
        :class:`FileMovedEvent`.
    """
    for root, directories, filenames in os.walk(dest_dir_path):
        for directory in directories:
            full_path = os.path.join(root, directory)
            renamed_path = full_path.replace(dest_dir_path, src_dir_path) if src_dir_path else None
            event = DirMovedEvent(renamed_path, full_path)
            event.is_synthetic = True
            yield event
        for filename in filenames:
            full_path = os.path.join(root, filename)
            renamed_path = full_path.replace(dest_dir_path, src_dir_path) if src_dir_path else None
            event = FileMovedEvent(renamed_path, full_path)
            event.is_synthetic = True
            yield event

def generate_sub_created_events(src_dir_path):
    """Generates an event list of :class:`DirCreatedEvent` and
    :class:`FileCreatedEvent` objects for all the files and directories within
    the given moved directory that were moved along with the directory.

    :param src_dir_path:
        The source path of the created directory.
    :returns:
        An iterable of file system events of type :class:`DirCreatedEvent` and
        :class:`FileCreatedEvent`.
    """
    for root, directories, filenames in os.walk(src_dir_path):
        for directory in directories:
            event = DirCreatedEvent(os.path.join(root, directory))
            event.is_synthetic = True
            yield event
        for filename in filenames:
            event = FileCreatedEvent(os.path.join(root, filename))
            event.is_synthetic = True
            yield event
<|endoftext|>


<commit_before>import math
import pdb
from .model import Model
from ...api import chain, clone, with_getitem, wrap, with_reshape
from .softmax import Softmax
from .layernorm import LayerNorm
from .resnet import Residual
from .affine import Affine
from .multiheaded_attention import MultiHeadedAttention

def with_reshape(layer):
    def with_reshape_forward(X, drop=0.):
        initial_shape = X.shape
        final_shape = list(initial_shape[:-1]) + [layer.nO]
        nB = X.shape[0]
        nL = X.shape[1]
        X2d = X.reshape(-1, X.shape[-1])
        X2d = X2d.astype(layer.ops.xp.float32)
        Y2d, Y2d_backprop = layer.begin_update(X2d, drop=drop)
        Y = Y2d.reshape(final_shape)

        def with_reshape_backward(dY, sgd=None):
            dY2d = dY.reshape(nB*nL, -1).astype(layer.ops.xp.float32)
            return Y2d_backprop(dY2d, sgd=sgd).reshape(initial_shape)
        return Y, with_reshape_backward
    return wrap(with_reshape_forward, layer)

class EncoderDecoder(Model):
    def __init__(self, nS=1, nH=6, nM=300, nTGT=10000):
        '''
        EncoderDecoder consists of an encoder stack, a decoder stack and an
        output layer which is a linear + softmax.
        Parameters explanation:
            nS: the number of encoders/decoders in the stack
            nH: the number of heads in the multiheaded attention
            nM: the token's embedding size
            nTGT: the number of unique words in output vocabulary
        '''
        Model.__init__(self)
        self.nS = nS
        self.nH = nH
        self.nM = nM
        self.nTGT = nTGT
        self.enc = clone(EncoderLayer(self.nH, self.nM), 1)
        # self.dec = clone(DecoderLayer(self.nH, self.nM), 1)
        self.dec = PoolingDecoder(self.nM)
        self.proj = with_reshape(Softmax(nO=nTGT, nI=nM))
        self._layers = [self.enc, self.dec, self.proj]

    def begin_update(self, inputs, drop=0.0):
        '''
        A batch object flows through the network. It contains input, output and
        corresponding masks. Input changes while the object travels through
        the network. Output is the golden output.
        Input: nB x nL x nM
        '''
        if len(inputs) == 2:
            (X0, Xmask), (Y0, Ymask) = inputs
            sentX = None
            sentY = None
        else:
            (X0, Xmask), (Y0, Ymask), (sentX, sentY) = inputs
        (X1, _), backprop_encode = self.enc.begin_update((X0, Xmask, sentX), drop=drop)
        (_, (Y1, _)), backprop_decode = self.dec.begin_update(((X1, Xmask, sentX), (Y0, Ymask, sentY)), drop=drop)
        word_probs, backprop_output = self.proj.begin_update(Y1, drop=drop)
        # Right-shift the word probabilities
        word_probs[:, 1:] = word_probs[:, :-1]
        word_probs[:, 0] = 0

        def finish_update(d_word_probs, sgd=None):
            # Unshift
            d_word_probs[:, :-1] = d_word_probs[:, 1:]
            d_word_probs[:, -1] = 0.

            dY1 = backprop_output(d_word_probs, sgd=sgd)
            zeros = Model.ops.xp.zeros(X0.shape, dtype=Model.ops.xp.float32)
            dX1, dY0 = backprop_decode((zeros, dY1), sgd=sgd)
            dX0 = backprop_encode(dX1, sgd=sgd)
            return (dX0, dY0)

        return (word_probs, Xmask), finish_update

def EncoderLayer(nH, nM):
    return chain(
        MultiHeadedAttention(nM, nH, layer='Encoder'),
        with_getitem(0, with_reshape(LayerNorm(Affine(nM, nM))))
    )

class PoolingDecoder(Model):
    def __init__(self, nM):
        Model.__init__(self)
        self.nM = nM
        self.ffd = LayerNorm(Maxout(nO=nM, nI=nM*3, pieces=3))
        self._layers = [self.ffd]

    def begin_update(self, X_Y, drop=0.):
        (X0, Xmask, _), (Y0, Ymask, _) = X_Y
        X_masked = self.ops.xp.copy(X0)
        X_masked[Xmask[:, 0, :] == 0] = -math.inf
        Xpool = X_masked.max(axis=1, keepdims=True)

        Y_masked = self.ops.xp.copy(Y0)
        Y_masked[Ymask[:, -1, :] == 0] = -math.inf
        Ypool = self.ops.allocate((X0.shape[0], X0.shape[1], self.nM))

        # maxing only over previous elements
        for i in range(X0.shape[1]):
            Ypool[:, i, :] = Y_masked[:, :i+1, :].max(axis=1, keepdims=True).squeeze()

        mixed = self.ops.allocate((X0.shape[0], X0.shape[1], self.nM*3))
        mixed[:, :, :self.nM] = Xpool
        mixed[:, :, self.nM:self.nM*2] = Ypool
        mixed[:, :, self.nM*2:] = Y0
        output, bp_output = self.ffd.begin_update(mixed.reshape((-1, self.nM*3)))
        output = output.reshape((X0.shape[0], X0.shape[1], output.shape[1]))

        def backprop_pooling_decoder(dX_d_output, sgd=None):
            dXin, d_output = dX_d_output
            d_output = d_output.reshape((X0.shape[0]*X0.shape[1], d_output.shape[2]))
            d_mixed = bp_output(d_output, sgd=sgd)
            d_mixed = d_mixed.reshape((X0.shape[0], X0.shape[1], d_mixed.shape[1]))
            dXpool = d_mixed[:, :, :self.nM]
            dYpool = d_mixed[:, :, self.nM:self.nM*2]
            dY0 = d_mixed[:, :, self.nM*2:]
            dX0 = self.ops.allocate(X0.shape)
            for i in range(X0.shape[0]):
                for j in range(X0.shape[1]):
                    for k in range(X0.shape[2]):
                        if X0[i, j, k] >= Xpool[i, 0, k]:
                            dX0[i, j, k] += dXpool[i, j, k]
            for i in range(Y0.shape[0]):
                for j in range(Y0.shape[1]):
                    for k in range(Y0.shape[2]):
                        if Y0[i, j, k] >= Ypool[i, j, k]:
                            dY0[i, j, k] += dYpool[i, j, k]
            return dXin + dX0, dY0
        return ((X0, Xmask), (output, Ymask)), backprop_pooling_decoder

class DecoderLayer(Model):
    def __init__(self, nH, nM):
        Model.__init__(self)
        self.nH = nH
        self.nM = nM
        self.x_attn = MultiHeadedAttention(nM, nH, layer='Decoder')
        self.y_attn = MultiHeadedAttention(nM, nH, layer='Decoder')
        self.ffd = with_reshape(LayerNorm(Affine(nM, nM)))
        self._layers = [self.x_attn, self.y_attn, self.ffd]

    def begin_update(self, X_Y, drop=0.0):
        (X0, Xmask, sentX), (Y0, Ymask, sentY) = X_Y
        (Y1, _), bp_self_attn = self.y_attn.begin_update((Y0, Ymask, sentY))
        (mixed, _), bp_mix_attn = self.x_attn.begin_update((Y1, X0, Xmask, sentY, sentX))
        output, bp_output = self.ffd.begin_update(mixed)

        def finish_update(dXprev_d_output, sgd=None):
            dXprev, d_output = dXprev_d_output
            d_mixed = bp_output(d_output, sgd=sgd)
            dY1, dX0 = bp_mix_attn(d_mixed, sgd=sgd)
            dY0 = bp_self_attn(dY1, sgd=sgd)
            return (dX0 + dXprev, dY0)

        return ((X0, Xmask), (output, Ymask)), finish_update
<commit_msg>Add correct encoder decoder implementation<commit_after># coding: utf8
from __future__ import unicode_literals, print_function
from .model import Model
from ...api import chain, clone, with_getitem, wrap, with_reshape
from .softmax import Softmax
from .relu import ReLu
from .layernorm import LayerNorm
from .maxout import Maxout
from .resnet import Residual
from .affine import Affine
from .multiheaded_attention import MultiHeadedAttention
from .positionwise_ffd import PositionwiseFeedForward
import copy
import math
import numpy as np
import torch.nn as nn
import torch
from thinc.extra.wrappers import PyTorchWrapper
import torch.nn.functional as F

class EncoderDecoder(Model):
    def __init__(self, nS=1, nH=6, nM=300, nTGT=10000, device='cpu'):
        '''
        EncoderDecoder consists of an encoder stack, a decoder stack and an
        output layer which is a linear + softmax.
        Parameters explanation:
            nS: the number of encoders/decoders in the stack
            nH: the number of heads in the multiheaded attention
            nM: the token's embedding size
            nTGT: the number of unique words in output vocabulary
        '''
        Model.__init__(self)
        self.nS = nS
        self.nH = nH
        self.nM = nM
        self.nTGT = nTGT
        self.device = device
        self.enc = Encoder(nM=nM, nH=nH, device=device, nS=nS)
        self.norm = PyTorchWrapper(PytorchLayerNorm(nM=nM, device=device))
        self.dec = clone(DecoderLayer(nM=nM, nH=nH, device=device), nS)
        self.proj = with_reshape(Softmax(nO=nTGT, nI=nM))
        self._layers = [self.enc, self.dec, self.proj]

    def begin_update(self, inputs, drop=0.1):
        '''
        A batch object flows through the network. It contains input, output and
        corresponding masks. Input changes while the object travels through
        the network. Output is the golden output.
        Input: nB x nL x nM
        '''
        X0, Xmask, Y0, Ymask = inputs
        X1, backprop_encode = self.enc.begin_update((X0, Xmask), drop=drop)
        (Y1, _, _, _), backprop_decode = self.dec.begin_update((Y0, X1, Xmask, Ymask), drop=drop)
        Y2, b_Y2 = self.norm.begin_update(Y1)
        word_probs, backprop_output = self.proj.begin_update(Y2, drop=drop)

        def finish_update(d_word_probs, sgd=None):
            dY2 = backprop_output(d_word_probs, sgd=sgd)
            dY1 = b_Y2(dY2, sgd=sgd)
            zeros = Model.ops.xp.zeros(X0.shape, dtype=Model.ops.xp.float32)
            dY0, dX1 = backprop_decode((dY1, zeros), sgd=sgd)
            dX0 = backprop_encode(dX1, sgd=sgd)
            return (dX0, dY0)
        return (word_probs, Xmask), finish_update

class PytorchLayerNorm(nn.Module):
    def __init__(self, nM=300, eps=1e-6, device='cpu'):
        super(PytorchLayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(nM).to(device))
        self.b_2 = nn.Parameter(torch.zeros(nM).to(device))
        self.eps = eps
        self.device = device

    def forward(self, x):
        mean = x.mean(-1, keepdim=True).to(self.device)
        std = x.std(-1, keepdim=True).to(self.device)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class Encoder(Model):
    def __init__(self, nM=300, nH=6, nS=6, device='cpu'):
        Model.__init__(self)
        self.stack = clone(EncoderLayer(nM=nM, nH=nH, device=device), nS)
        self.norm = PyTorchWrapper(PytorchLayerNorm(nM=nM, device=device))

    def begin_update(self, input, drop=0.1):
        X0, mask = input
        (X1, _), b_X1 = self.stack.begin_update((X0, mask), drop=0.1)
        X2, b_X2 = self.norm.begin_update(X1)

        def finish_update(dX2, sgd=None):
            dX1 = b_X2(dX2, sgd=sgd)
            dX0 = b_X1(dX1, sgd=sgd)
            return dX0
        return X2, finish_update

class EncoderLayer(Model):
    def __init__(self, nM=300, nH=6, device='cpu'):
        Model.__init__(self)
        self.attn = MultiHeadedAttention(nM=nM, nH=nH)
        self.ffd = PositionwiseFeedForward(nM, 4 * nM)
        self.norm = PyTorchWrapper(PytorchLayerNorm(nM, device=device))
        self.nM = nM
        self.layers_ = [self.attn, self.ffd, self.norm]

    def begin_update(self, input, drop=0.1):
        X0, mask = input
        X1, b_X1 = self.attn.begin_update((X0, mask, None), drop=drop)
        X2, b_X2 = self.norm.begin_update(X1)
        X3 = X0 + X2

        X4, b_X4 = self.ffd.begin_update(X3, drop=drop)
        X5, b_X5 = self.norm.begin_update(X4)
        X6 = X3 + X5

        def finish_update(dX6, sgd=None):
            dX5 = dX6
            dX4 = b_X5(dX5, sgd=sgd)
            dX3 = b_X4(dX4, sgd=sgd)
            dX3 += dX6

            dX2 = dX3
            dX1 = b_X2(dX2, sgd=sgd)
            dX0 = b_X1(dX1, sgd=sgd)

            dX0 += dX3
            return X0
        return (X6, mask), finish_update

class DecoderLayer(Model):
    def __init__(self, nM=300, nH=6, device='cpu'):
        Model.__init__(self)
        self.y_attn = MultiHeadedAttention(nM=nM, nH=nH)
        self.x_attn = MultiHeadedAttention(nM=nM, nH=nH)
        self.norm = PyTorchWrapper(PytorchLayerNorm(nM, device=device))
        self.ffd = PositionwiseFeedForward(nM, 4 * nM)
        self.layers_ = [self.norm, self.y_attn, self.x_attn, self.ffd]

    def begin_update(self, input, drop=0.1):
        Y0, X0, X_mask, Y_mask = input
        Y1, b_Y1 = self.y_attn.begin_update((Y0, Y_mask, None), drop=drop)
        Y2, b_Y2 = self.norm.begin_update(Y1)
        Y3 = Y0 + Y2
        Y4, b_Y4 = self.x_attn.begin_update((Y3, X0, X_mask, None, None), drop=drop)
        Y5, b_Y5 = self.norm.begin_update(Y4)
        Y6 = Y3 + Y5
        Y7, b_Y7 = self.ffd.begin_update(Y6, drop=drop)

        def finish_update(dI, sgd=None):
            dY7, dX = dI
            dY6 = b_Y7(dY7, sgd=sgd)
            dY5 = dY6
            dY4 = b_Y5(dY5, sgd=sgd)
            dY3, dX0 = b_Y4(dY4, sgd=sgd)
            dY3 += dY6
            dY2 = dY3
            dY1 = b_Y2(dY2, sgd=sgd)
            dY0 = b_Y1(dY1, sgd=sgd)
            dY0 += dY3
            dX0 += dX
            return (dY0, dX0)
        return (Y7, X0, X_mask, Y_mask), finish_update
<|endoftext|>


<commit_before># coding=utf-8
import commands
from copy_reg import pickle
from types import MethodType
from multiprocessing import Pool

from settings import max_running_number, lrun_gid, lrun_uid, use_tmpfs
from consts import Language, Result

# 下面两个函数告诉Python怎么pickle类实例中的方法，否则Python2会报错，是Python2的已知bug
# http://stackoverflow.com/questions/1816958/cant-pickle-type-instancemethod-when-using-pythons-multiprocessing-pool-ma/7309686
def _pickle_method(method):
    func_name = method.im_func.__name__
    obj = method.im_self
    cls = method.im_class
    return _unpickle_method, (func_name, obj, cls)

def _unpickle_method(func_name, obj, cls):
    for cls in cls.mro():
        try:
            func = cls.__dict__[func_name]
        except KeyError:
            pass
        else:
            break
    return func.__get__(obj, cls)

class JudgeClientException(Exception):
    pass

class JudgeClient(object):
    def __init__(self, language, exec_file_path, max_cpu_time,
                 max_real_time, max_memory, test_case_dir):
        """
        :param language: 语言，见consts.py
        :param exec_file_path: 可执行文件路径
        :param max_cpu_time: 最大cpu时间，单位ms
        :param max_real_time: 最大执行时间，单位ms
        :param max_memory: 最大内存，单位MB
        :param test_case_dir: 测试用户文件夹路径
        :return:返回结果list
        """
        self.language = language
        self.exec_file_path = exec_file_path
        self.max_cpu_time = max_cpu_time
        self.max_real_time = max_real_time
        self.max_memory = max_memory
        self.test_case_dir = test_case_dir
        # 进程池
        self.pool = Pool(processes=max_running_number)
        # 结果数组
        self.results = []
        # 测试用例配置项
        self.test_case_info = self.load_test_case_info()

    def load_test_case_info(self):
        # 读取测试用例信息 转换为dict
        # try:
        #     f = open(self.test_case_dir + "info")
        #     return json.loads(f.read())
        # except IOError:
        #     raise JudgeClientException("Test case config file not found")
        # except ValueError:
        #     raise JudgeClientException("Test case config file format error")
        return {"test_case_number": 2,
                "test_cases":
                    {
                        "1": {"input_name": "1.in",
                              "output_name": "1.out",
                              "output_md5": "yyy",
                              "output_size": 100},

                        "2": {"input_name": "2.in",
                              "output_name": "2.out",
                              "output_md5": "yyy",
                              "output_size": 100}
                    }
                }

    def generate_command(self, test_case_id):
        """
        设置相关运行限制 进制访问网络 如果启用tmpfs 就把代码输出写入tmpfs，否则写入硬盘
        """
        # todo 系统调用白名单 chroot等参数
        # fixme 时间的单位问题
        command = "lrun" + \
                  " --max-cpu-time " + str(self.max_cpu_time / 1000.0) + \
                  " --max-real-time " + str(self.max_real_time / 1000.0) + \
                  " --max-memory " + str(self.max_memory * 1000 * 1000) + \
                  " --network false" + \
                  " --uid " + str(lrun_uid) + \
                  " --gid " + str(lrun_gid)
        #if use_tmpfs:
        #    command += (" --tmpfs /var " +
        #                str(int(self.test_case_info["test_cases"][str(test_case_id)]["output_size"] * 1.2)))

        if self.language == Language.JAVA:
            command += (" java " + self.exec_file_path)
        else:
            command += (" " + self.exec_file_path)
        # fixme 输出路径
        command += (" 0<" + self.test_case_dir + str(test_case_id) + ".in" +
                    " 1>" + "/var/judge/" + str(test_case_id) + ".out" +
                    " 3>&2")
        return command

    def parse_lrun_output(self, output):
        lines = output.split("\n")
        if len(lines) != 7:
            raise JudgeClientException("Lrun result parse error")
        result = {}
        # 将lrun输出的各种带下划线 不带下划线的字符串统一处理
        translate = {"MEMORY": "memory",
                     "CPUTIME": "cpu_time",
                     "CPU_TIME": "cpu_time",
                     "REALTIME": "real_time",
                     "REAL_TIME": "real_time",
                     "TERMSIG": "term_sig",
                     "SIGNALED": "siginaled",
                     "EXITCODE": "exit_code",
                     "EXCEED": "exceed"}
        for line in lines:
            name = line[:9].strip(" ")
            value = line[9:]
            if name == "MEMORY":
                result[translate[name]] = int(value)
            elif name == "CPUTIME":
                result[translate[name]] = float(value) * 1000
            elif name == "REALTIME":
                result[translate[name]] = float(value) * 1000
            elif name == "EXITCODE":
                result[translate[name]] = int(value)
            elif name == "TERMSIG":
                result[translate[name]] = int(value)
            elif name == "SIGNALED":
                result[translate[name]] = int(value)
            elif name == "EXCEED":
                if value == "none":
                    result[translate[name]] = None
                else:
                    result[translate[name]] = translate[value]
        return result

    def judge_one(self, test_case_id):
        # 运行lrun程序 接收返回值
        command = self.generate_command(test_case_id)
        status_code, output = commands.getstatusoutput(command)
        if status_code:
            raise JudgeClientException(output)
        run_result = self.parse_lrun_output(output)

        run_result["test_case_id"] = test_case_id

        # 如果返回值非0 或者信号量不是0 代表非正常结束
        if run_result["exit_code"] or run_result["term_sig"] or run_result["siginaled"]:
            run_result["result"] = Result.RUNTIME_ERROR
            return run_result

        # 代表内存或者时间超过限制了
        if run_result["exceed"]:
            if run_result["exceed"] == "memory":
                run_result["result"] = Result.MEMORY_LIMIT_EXCEEDED
            elif run_result["exceed"] in ["cpu_time", "real_time"]:
                run_result["result"] = Result.TIME_LIMIT_EXCEEDED
            else:
                raise JudgeClientException("Error exceeded type: " + run_result["exceed"])
            return run_result

        # 下面就是代码正常运行了
        run_result["result"] = Result.ACCEPTED
        return run_result

    def collect_result(self, result):
        self.results.append(result)

    def run(self):
        # 添加到任务队列
        _results = []
        results = []
        for i in range(self.test_case_info["test_case_number"]):
            _results.append(self.pool.apply_async(self.judge_one, args=(i + 1, )))
        self.pool.close()
        self.pool.join()
        for item in _results:
            # 注意多进程中的异常只有在get()的时候才会被引发
            # http://stackoverflow.com/questions/22094852/how-to-catch-exceptions-in-workers-in-multiprocessing
            try:
                results.append(item.get())
            except Exception as e:
                results.append({"result": Result.SYSTEM_ERROR})
        return results

    def __getstate__(self):
        # 不同的pool之间进行pickle的时候要排除自己，否则报错
        # http://stackoverflow.com/questions/25382455/python-notimplementederror-pool-objects-cannot-be-passed-between-processes
        self_dict = self.__dict__.copy()
        del self_dict['pool']
        return self_dict

pickle(MethodType, _pickle_method, _unpickle_method)
client = JudgeClient(language=Language.C,
                     exec_file_path="/var/judge/a.out",
                     max_cpu_time=1000000,
                     max_real_time=200000,
                     max_memory=1,
                     test_case_dir="/var/test_case/1/")
print client.run()
<commit_msg>使用更好的办法修复Python2不能pickle实例方法的问题<commit_after># coding=utf-8
import commands
from copy_reg import pickle
from types import MethodType
from multiprocessing import Pool

from settings import max_running_number, lrun_gid, lrun_uid, use_tmpfs
from consts import Language, Result

class JudgeClientException(Exception):
    pass

# 下面这个函数作为代理访问实例变量，否则Python2会报错，是Python2的已知问题
# http://stackoverflow.com/questions/1816958/cant-pickle-type-instancemethod-when-using-pythons-multiprocessing-pool-ma/7309686
def _run(instance, test_case_id):
    return instance.judge_one(test_case_id)

class JudgeClient(object):
    def __init__(self, language, exec_file_path, max_cpu_time,
                 max_real_time, max_memory, test_case_dir):
        """
        :param language: 语言，见consts.py
        :param exec_file_path: 可执行文件路径
        :param max_cpu_time: 最大cpu时间，单位ms
        :param max_real_time: 最大执行时间，单位ms
        :param max_memory: 最大内存，单位MB
        :param test_case_dir: 测试用户文件夹路径
        :return:返回结果list
        """
        self.language = language
        self.exec_file_path = exec_file_path
        self.max_cpu_time = max_cpu_time
        self.max_real_time = max_real_time
        self.max_memory = max_memory
        self.test_case_dir = test_case_dir
        # 进程池
        self.pool = Pool(processes=max_running_number)
        # 结果数组
        self.results = []
        # 测试用例配置项
        self.test_case_info = self.load_test_case_info()

    def load_test_case_info(self):
        # 读取测试用例信息 转换为dict
        # try:
        #     f = open(self.test_case_dir + "info")
        #     return json.loads(f.read())
        # except IOError:
        #     raise JudgeClientException("Test case config file not found")
        # except ValueError:
        #     raise JudgeClientException("Test case config file format error")
        return {"test_case_number": 2,
                "test_cases":
                    {
                        "1": {"input_name": "1.in",
                              "output_name": "1.out",
                              "output_md5": "yyy",
                              "output_size": 100},

                        "2": {"input_name": "2.in",
                              "output_name": "2.out",
                              "output_md5": "yyy",
                              "output_size": 100}
                    }
                }

    def generate_command(self, test_case_id):
        """
        设置相关运行限制 进制访问网络 如果启用tmpfs 就把代码输出写入tmpfs，否则写入硬盘
        """
        # todo 系统调用白名单 chroot等参数
        # fixme 时间的单位问题
        command = "lrun" + \
                  " --max-cpu-time " + str(self.max_cpu_time / 1000.0) + \
                  " --max-real-time " + str(self.max_real_time / 1000.0) + \
                  " --max-memory " + str(self.max_memory * 1000 * 1000) + \
                  " --network false" + \
                  " --uid " + str(lrun_uid) + \
                  " --gid " + str(lrun_gid)
        #if use_tmpfs:
        #    command += (" --tmpfs /var " +
        #                str(int(self.test_case_info["test_cases"][str(test_case_id)]["output_size"] * 1.2)))

        if self.language == Language.JAVA:
            command += (" java " + self.exec_file_path)
        else:
            command += (" " + self.exec_file_path)
        # fixme 输出路径
        command += (" 0<" + self.test_case_dir + str(test_case_id) + ".in" +
                    " 1>" + "/var/judge/" + str(test_case_id) + ".out" +
                    " 3>&2")
        return command

    def parse_lrun_output(self, output):
        lines = output.split("\n")
        if len(lines) != 7:
            raise JudgeClientException("Lrun result parse error")
        result = {}
        # 将lrun输出的各种带下划线 不带下划线的字符串统一处理
        translate = {"MEMORY": "memory",
                     "CPUTIME": "cpu_time",
                     "CPU_TIME": "cpu_time",
                     "REALTIME": "real_time",
                     "REAL_TIME": "real_time",
                     "TERMSIG": "term_sig",
                     "SIGNALED": "siginaled",
                     "EXITCODE": "exit_code",
                     "EXCEED": "exceed"}
        for line in lines:
            name = line[:9].strip(" ")
            value = line[9:]
            if name == "MEMORY":
                result[translate[name]] = int(value)
            elif name == "CPUTIME":
                result[translate[name]] = float(value) * 1000
            elif name == "REALTIME":
                result[translate[name]] = float(value) * 1000
            elif name == "EXITCODE":
                result[translate[name]] = int(value)
            elif name == "TERMSIG":
                result[translate[name]] = int(value)
            elif name == "SIGNALED":
                result[translate[name]] = int(value)
            elif name == "EXCEED":
                if value == "none":
                    result[translate[name]] = None
                else:
                    result[translate[name]] = translate[value]
        return result

    def judge_one(self, test_case_id):
        # 运行lrun程序 接收返回值
        command = self.generate_command(test_case_id)
        status_code, output = commands.getstatusoutput(command)
        if status_code:
            raise JudgeClientException(output)
        run_result = self.parse_lrun_output(output)

        run_result["test_case_id"] = test_case_id

        # 如果返回值非0 或者信号量不是0 代表非正常结束
        if run_result["exit_code"] or run_result["term_sig"] or run_result["siginaled"]:
            run_result["result"] = Result.RUNTIME_ERROR
            return run_result

        # 代表内存或者时间超过限制了
        if run_result["exceed"]:
            if run_result["exceed"] == "memory":
                run_result["result"] = Result.MEMORY_LIMIT_EXCEEDED
            elif run_result["exceed"] in ["cpu_time", "real_time"]:
                run_result["result"] = Result.TIME_LIMIT_EXCEEDED
            else:
                raise JudgeClientException("Error exceeded type: " + run_result["exceed"])
            return run_result

        # 下面就是代码正常运行了
        run_result["result"] = Result.ACCEPTED
        return run_result

    def collect_result(self, result):
        self.results.append(result)

    def run(self):
        # 添加到任务队列
        _results = []
        results = []
        for i in range(self.test_case_info["test_case_number"]):
            _results.append(self.pool.apply_async(_run, (self, i + 1)))
        self.pool.close()
        self.pool.join()
        for item in _results:
            # 注意多进程中的异常只有在get()的时候才会被引发
            # http://stackoverflow.com/questions/22094852/how-to-catch-exceptions-in-workers-in-multiprocessing
            try:
                results.append(item.get())
            except Exception as e:
                results.append({"result": Result.SYSTEM_ERROR})
        return results

    def __getstate__(self):
        # 不同的pool之间进行pickle的时候要排除自己，否则报错
        # http://stackoverflow.com/questions/25382455/python-notimplementederror-pool-objects-cannot-be-passed-between-processes
        self_dict = self.__dict__.copy()
        del self_dict['pool']
        return self_dict

# pickle(MethodType, _pickle_method, _unpickle_method)
client = JudgeClient(language=Language.C,
                     exec_file_path="/var/judge/a.out",
                     max_cpu_time=1000000,
                     max_real_time=200000,
                     max_memory=1,
                     test_case_dir="/var/test_case/1/")
print client.run()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import logging
import threading
import uuid

import psycopg2
import psycopg2.extras
import psycopg2.extensions

from jukoro.decorators import raise_if
from jukoro.structures import LockRing

from jukoro.pg.exceptions import (
    PgPoolClosedError, PgConnectionClosedError, PgCursorClosedError,
    PgDoesNotExistError)
from jukoro.pg.utils import pg_uri_to_kwargs

# block size for PostgreSQL fetched rows iteration
BLOCK_SIZE = 2048

# module level logger
logger = logging.getLogger(__name__)
# sql queries logger
sql_logger = logging.getLogger('jukoro.pg.sql')

def is_closed(instance, *args, **kwargs):
    """
    Helper function to test if some instance is closed (to be used with
    ``raise_if`` decorator)

    :param instance: some instance with ``is_closed`` property accessible
    :returns:        value of ``instance.is_closed`` or False

    """
    return getattr(instance, 'is_closed', False)

# decorator to test if cursor closed
raise_if_cursor_closed = raise_if(PgCursorClosedError,
                                  'cursor closed', is_closed)
# decorator to test if connection closed
raise_if_connection_closed = raise_if(PgConnectionClosedError,
                                      'connection closed', is_closed)
# decorator to test if pool closed
raise_if_pool_closed = raise_if(PgPoolClosedError,
                                'pool closed', is_closed)

class PgResult(object):
    """
    Provides methods to work with query results

    :param cursor: instance of `psycopg2.extensions.cursor`

    """

    __slots__ = ('_cursor', )

    def __init__(self, cursor):
        self._cursor = cursor

    @property
    def is_closed(self):
        """
        Returns current status of the instance

        """
        return self._cursor is None

    @raise_if_cursor_closed
    def get(self):
        """
        Method to get first row from query result

        :returns: fetched from db row
        :rtype:   dict (`psycopg2.extras.RealDictCursor`)
        :raises PgDoesNotExistError: if query returned no results

        """
        try:
            return self[0]
        except psycopg2.ProgrammingError:
            raise PgDoesNotExistError

    @raise_if_cursor_closed
    def all(self):
        """
        Method to fetch all rows from cursor

        :returns: all rows fetched from db
        :rtype:   list

        """
        return self._cursor.fetchall()

    @raise_if_cursor_closed
    def block(self):
        """
        Method to fetch block of rows from cursor

        :returns: block of rows fetched from db
        :rtype:   list

        """
        return self._cursor.fetchmany()

    def __iter__(self):
        """
        Generator function to iterate over the rows from db
        (iterates using blocks to keep memory usage low if possible)

        """
        # named cursor will transparently work the same way the client-side
        # cursor does
        block = self.block()
        while block:
            for it in block:
                yield it
            block = self.block()

    @property
    @raise_if_cursor_closed
    def rowcount(self):
        """
        Returns `psycopg2.extensions.cursor.rowcount` read-only attribute value

        """
        return self._cursor.rowcount

    @raise_if_cursor_closed
    def __len__(self):
        """
        Magic method to get length of the instance

        Returns number of rows fetched or affected by query
        In case cursor has a name (ie it's a named cursor) return 0

        """
        if self._cursor.name is not None:
            return 0
        return self._cursor.rowcount

    def __getitem__(self, value):
        """
        Magic method to provide access to specific row in results or
        to slice query results

        :param value: index (integer) or slice
        :returns:     single row or list of rows

        NB. doesn't support step value for slicing

        """
        if isinstance(value, int):
            self.scroll(value)
            resp = self._cursor.fetchone()
            if resp:
                return resp
            return
        if isinstance(value, slice):
            # TODO raise for negative positions
            start = value.start or 0
            stop = value.stop
            if stop is None:
                raise ValueError('undefined value for upper boundary')
            # TODO respect step
            # step = value.step
            self.scroll(start)
            return self._cursor.fetchmany(stop - start)

    def close(self):
        """
        Method to close results after processing (free cursor instance)

        """
        self._cursor = None

    @raise_if_cursor_closed
    def scroll(self, pos):
        """
        Method to scroll to specific position within fetched results

        :param pos: (int) position to scroll to
        :raises PgDoesNotExistError: if position is unavailable

        """
        try:
            self._cursor.scroll(pos, 'absolute')
        except (psycopg2.ProgrammingError, IndexError):
            raise PgDoesNotExistError

class PgTransaction(object):

    __slots__ = ('_pg_conn', '_autocommit', '_named', '_cursor', '_failed',
                 '_result', '_closed', '_queries', '_block_size')

    def __init__(self, conn, autocommit=True, named=False, **kwargs):
        if named and autocommit:
            logger.warn(
                'incompatible parameters "autocommit = named = True"')
            autocommit = False
        self._pg_conn = conn  # instance of PgConnection
        self._autocommit = autocommit
        self._named = named
        self._cursor = None
        self._failed = False
        self._result = None
        self._closed = False
        self._queries = []
        self._block_size = kwargs.get('block_size', BLOCK_SIZE)

    def _ensure_cursor(self):
        if self._cursor is not None:
            return
        if self._named or not self._autocommit:
            if self._pg_conn.autocommit:
                self._pg_conn.autocommit = False
        else:
            if not self._pg_conn.autocommit:
                self._pg_conn.autocommit = True
        self._cursor = self._pg_conn.cursor(self._named)
        self._cursor.arraysize = self._block_size
        if self._named:
            self._cursor.itersize = self._block_size

    def __enter__(self):
        self._ensure_cursor()
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is not None:
            logger.exception('exception executing query')
            if not self._autocommit:
                self._pg_conn.rollback()
            self._failed = True
        else:
            if self._named or not self._autocommit:
                self._pg_conn.commit()
        self.close()

    @property
    def block_size(self):
        return self._block_size

    @property
    def arraysize(self):
        self._ensure_cursor()
        return self._cursor.arraysize

    @property
    def itersize(self):
        self._ensure_cursor()
        return self._cursor.itersize

    @property
    def is_closed(self):
        return self._closed

    @property
    def queries(self):
        for q in self._queries:
            yield q

    def _should_close_cursor(self):
        if (not self._named and self._cursor is not None) \
                or (self._named and not self._failed and self._queries):
            return True
        return False

    def close(self):
        self._close_result()
        if self._should_close_cursor():
            self._cursor.close()
        self._cursor = self._queries = None
        self._pg_conn.reattach()
        self._pg_conn = None
        self._closed = True

    def _close_result(self):
        if self._result is not None:
            self._result.close()
            self._result = None

    def execute(self, query, params=None):
        return self._exec(query, params=params)

    def execute_and_get(self, query, params=None):
        return self.execute(query, params).get()

    def callproc(self, procname, params=None):
        return self._exec(procname, params=params, proc=True)

    def _exec(self, q_or_proc, params=None, proc=False):
        self._ensure_cursor()
        fn = self._cursor.callproc if proc else self._cursor.execute
        sql_logger.debug(
            'executing query "%s" with "%s"', q_or_proc, params or [])
        fn(q_or_proc, params)
        self._queries.append(self._cursor.query)
        self._close_result()
        self._result = PgResult(self._cursor)
        return self._result

class PgConnection(object):

    __slots__ = ('_uri', '_schema', '_pg_pool', '_conn_kwargs', '_conn',
                 '_autoclose', '_closed')

    def __init__(self, uri, pool=None, autoclose=False):
        self._uri = uri
        self._pg_pool = pool
        kwargs = pg_uri_to_kwargs(uri)
        self._schema = kwargs.pop('schema')
        self._conn_kwargs = kwargs
        self._conn = None  # psycopg2.connection
        self._autoclose = autoclose
        self._closed = False

        # logger.debug('connection created %s', repr(self))

    def __repr__(self):
        return '<PgConnection(uri="{}")> at {}'.format(self._uri,
                                                       hex(id(self)))

    @property
    @raise_if_connection_closed
    def conn(self):
        if self._conn is None:
            self._conn = _connect(**self._conn_kwargs)
            self._conn.set_session(
                psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED,
                autocommit=True)
            self._conn.set_client_encoding('UTF8')
            self._conn.autocommit = True
            cur = self._conn.cursor()
            cur.execute('SET TIME ZONE "UTC";')
            cur.execute('SET search_path TO %s;', (self.schema, ))
            cur.close()
        return self._conn

    @property
    def autocommit(self):
        return self.conn.autocommit

    @autocommit.setter
    def autocommit(self, value):
        self.conn.autocommit = value

    @property
    def is_closed(self):
        return self._closed

    @property
    def schema(self):
        return self._schema

    @raise_if_connection_closed
    def commit(self):
        self.conn.commit()

    @raise_if_connection_closed
    def rollback(self):
        self.conn.rollback()

    @raise_if_connection_closed
    def close(self):
        if self._conn is not None:
            self._conn.close()
            self._conn = None
        self._pg_pool = None
        self._closed = True

    @raise_if_connection_closed
    def cursor(self, named=False):
        if named:
            return self.conn.cursor(
                name=str(uuid.uuid4()), scrollable=True, withhold=True)
        return self.conn.cursor()

    @raise_if_connection_closed
    def reattach(self):
        if self._pg_pool is not None:
            self._pg_pool.unlock(self)
        elif self._autoclose:
            self.close()

    @raise_if_connection_closed
    def transaction(self, **kwargs):
        return PgTransaction(self, **kwargs)

class PgDbPool(object):

    __slots__ = ('_uri', '_pool_size', '_pool',
                 '_warmed_up', '_closed', '_lock')

    def __init__(self, uri, pool_size=5, **kwargs):
        self._uri = uri
        self._pool_size = pool_size
        self._pool = LockRing()
        self._warmed_up = False
        self._closed = False
        self._lock = threading.Lock()

    @property
    def is_closed(self):
        return self._closed

    @property
    def uri(self):
        return self._uri

    def __repr__(self):
        return '<PgDbPool("{}")> at {}'.format(self._uri, hex(id(self)))

    def __len__(self):
        return len(self._pool)

    def close(self):
        with self._lock:
            for __ in xrange(len(self._pool)):
                conn = self._pool.pop()
                conn.close()
            self._closed = True
            self._pool.reset()

    @raise_if_pool_closed
    def transaction(self, **kwargs):
        with self._lock:
            conn = self._get_conn()
        return conn.transaction(**kwargs)

    def _get_conn(self):
        if not self._warmed_up:
            self._warm_up()
        try:
            conn = self._pool.next()
        except (IndexError, StopIteration):
            logger.error('pool exhausted, making new connection')
            conn = self._new_conn(autoclose=True)
        return conn

    def _new_conn(self, **kwargs):
        return PgConnection(self._uri, **kwargs)

    @raise_if_pool_closed
    def unlock(self, conn):
        self._pool.push(conn)

    def _warm_up(self):
        self._warmed_up = True
        for __ in xrange(self._pool_size):
            self._pool.push(self._new_conn(pool=self))
        logger.info(
            'warmed up pool "%s" length %s', repr(self), self._pool_size)

def _connect(**kwargs):
    kwargs.setdefault('cursor_factory', psycopg2.extras.RealDictCursor)
    return psycopg2.connect(**kwargs)
<commit_msg>module docstring<commit_after># -*- coding: utf-8 -*-
"""
Module to provide abstractions to work with PostgreSQL built on top of
``psycopg2.extensions.connection`` and ``psycopg2.extensions.cursor``.

- :class:`PgDbPool <jukoro.pg.db.PgDbPool>` - abstraction for a pool of
connections
- :class:`PgConnection <jukoro.pg.db.PgConnection>` - abstraction for a
connection
- :class:`PgTransaction <jukoro.pg.db.PgTransaction>` - abstraction for a
transaction
- :class:`PgResult <jukoro.pg.db.PgResult>` - abstraction for results

"""

import logging
import threading
import uuid

import psycopg2
import psycopg2.extras
import psycopg2.extensions

from jukoro.decorators import raise_if
from jukoro.structures import LockRing

from jukoro.pg.exceptions import (
    PgPoolClosedError, PgConnectionClosedError, PgCursorClosedError,
    PgDoesNotExistError)
from jukoro.pg.utils import pg_uri_to_kwargs

# block size for PostgreSQL fetched rows iteration
BLOCK_SIZE = 2048

# module level logger
logger = logging.getLogger(__name__)
# sql queries logger
sql_logger = logging.getLogger('jukoro.pg.sql')

def is_closed(instance, *args, **kwargs):
    """
    Helper function to test if some instance is closed (to be used with
    ``raise_if`` decorator)

    :param instance: some instance with ``is_closed`` property accessible
    :returns:        value of ``instance.is_closed`` or False

    """
    return getattr(instance, 'is_closed', False)

# decorator to test if cursor closed
raise_if_cursor_closed = raise_if(PgCursorClosedError,
                                  'cursor closed', is_closed)
# decorator to test if connection closed
raise_if_connection_closed = raise_if(PgConnectionClosedError,
                                      'connection closed', is_closed)
# decorator to test if pool closed
raise_if_pool_closed = raise_if(PgPoolClosedError,
                                'pool closed', is_closed)

class PgResult(object):
    """
    Provides methods to work with query results

    :param cursor: instance of `psycopg2.extensions.cursor`

    """

    __slots__ = ('_cursor', )

    def __init__(self, cursor):
        self._cursor = cursor

    @property
    def is_closed(self):
        """
        Returns current status of the instance

        """
        return self._cursor is None

    @raise_if_cursor_closed
    def get(self):
        """
        Method to get first row from query result

        :returns: fetched from db row
        :rtype:   dict (`psycopg2.extras.RealDictCursor`)
        :raises PgDoesNotExistError: if query returned no results

        """
        try:
            return self[0]
        except psycopg2.ProgrammingError:
            raise PgDoesNotExistError

    @raise_if_cursor_closed
    def all(self):
        """
        Method to fetch all rows from cursor

        :returns: all rows fetched from db
        :rtype:   list

        """
        return self._cursor.fetchall()

    @raise_if_cursor_closed
    def block(self):
        """
        Method to fetch block of rows from cursor

        :returns: block of rows fetched from db
        :rtype:   list

        """
        return self._cursor.fetchmany()

    def __iter__(self):
        """
        Generator function to iterate over the rows from db
        (iterates using blocks to keep memory usage low if possible)

        """
        # named cursor will transparently work the same way the client-side
        # cursor does
        block = self.block()
        while block:
            for it in block:
                yield it
            block = self.block()

    @property
    @raise_if_cursor_closed
    def rowcount(self):
        """
        Returns `psycopg2.extensions.cursor.rowcount` read-only attribute value

        """
        return self._cursor.rowcount

    @raise_if_cursor_closed
    def __len__(self):
        """
        Magic method to get length of the instance

        Returns number of rows fetched or affected by query
        In case cursor has a name (ie it's a named cursor) return 0

        """
        if self._cursor.name is not None:
            return 0
        return self._cursor.rowcount

    def __getitem__(self, value):
        """
        Magic method to provide access to specific row in results or
        to slice query results

        :param value: index (integer) or slice
        :returns:     single row or list of rows

        NB. doesn't support step value for slicing

        """
        if isinstance(value, int):
            self.scroll(value)
            resp = self._cursor.fetchone()
            if resp:
                return resp
            return
        if isinstance(value, slice):
            # TODO raise for negative positions
            start = value.start or 0
            stop = value.stop
            if stop is None:
                raise ValueError('undefined value for upper boundary')
            # TODO respect step
            # step = value.step
            self.scroll(start)
            return self._cursor.fetchmany(stop - start)

    def close(self):
        """
        Method to close results after processing (free cursor instance)

        """
        self._cursor = None

    @raise_if_cursor_closed
    def scroll(self, pos):
        """
        Method to scroll to specific position within fetched results

        :param pos: (int) position to scroll to
        :raises PgDoesNotExistError: if position is unavailable

        """
        try:
            self._cursor.scroll(pos, 'absolute')
        except (psycopg2.ProgrammingError, IndexError):
            raise PgDoesNotExistError

class PgTransaction(object):

    __slots__ = ('_pg_conn', '_autocommit', '_named', '_cursor', '_failed',
                 '_result', '_closed', '_queries', '_block_size')

    def __init__(self, conn, autocommit=True, named=False, **kwargs):
        if named and autocommit:
            logger.warn(
                'incompatible parameters "autocommit = named = True"')
            autocommit = False
        self._pg_conn = conn  # instance of PgConnection
        self._autocommit = autocommit
        self._named = named
        self._cursor = None
        self._failed = False
        self._result = None
        self._closed = False
        self._queries = []
        self._block_size = kwargs.get('block_size', BLOCK_SIZE)

    def _ensure_cursor(self):
        if self._cursor is not None:
            return
        if self._named or not self._autocommit:
            if self._pg_conn.autocommit:
                self._pg_conn.autocommit = False
        else:
            if not self._pg_conn.autocommit:
                self._pg_conn.autocommit = True
        self._cursor = self._pg_conn.cursor(self._named)
        self._cursor.arraysize = self._block_size
        if self._named:
            self._cursor.itersize = self._block_size

    def __enter__(self):
        self._ensure_cursor()
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is not None:
            logger.exception('exception executing query')
            if not self._autocommit:
                self._pg_conn.rollback()
            self._failed = True
        else:
            if self._named or not self._autocommit:
                self._pg_conn.commit()
        self.close()

    @property
    def block_size(self):
        return self._block_size

    @property
    def arraysize(self):
        self._ensure_cursor()
        return self._cursor.arraysize

    @property
    def itersize(self):
        self._ensure_cursor()
        return self._cursor.itersize

    @property
    def is_closed(self):
        return self._closed

    @property
    def queries(self):
        for q in self._queries:
            yield q

    def _should_close_cursor(self):
        if (not self._named and self._cursor is not None) \
                or (self._named and not self._failed and self._queries):
            return True
        return False

    def close(self):
        self._close_result()
        if self._should_close_cursor():
            self._cursor.close()
        self._cursor = self._queries = None
        self._pg_conn.reattach()
        self._pg_conn = None
        self._closed = True

    def _close_result(self):
        if self._result is not None:
            self._result.close()
            self._result = None

    def execute(self, query, params=None):
        return self._exec(query, params=params)

    def execute_and_get(self, query, params=None):
        return self.execute(query, params).get()

    def callproc(self, procname, params=None):
        return self._exec(procname, params=params, proc=True)

    def _exec(self, q_or_proc, params=None, proc=False):
        self._ensure_cursor()
        fn = self._cursor.callproc if proc else self._cursor.execute
        sql_logger.debug(
            'executing query "%s" with "%s"', q_or_proc, params or [])
        fn(q_or_proc, params)
        self._queries.append(self._cursor.query)
        self._close_result()
        self._result = PgResult(self._cursor)
        return self._result

class PgConnection(object):

    __slots__ = ('_uri', '_schema', '_pg_pool', '_conn_kwargs', '_conn',
                 '_autoclose', '_closed')

    def __init__(self, uri, pool=None, autoclose=False):
        self._uri = uri
        self._pg_pool = pool
        kwargs = pg_uri_to_kwargs(uri)
        self._schema = kwargs.pop('schema')
        self._conn_kwargs = kwargs
        self._conn = None  # psycopg2.connection
        self._autoclose = autoclose
        self._closed = False

        # logger.debug('connection created %s', repr(self))

    def __repr__(self):
        return '<PgConnection(uri="{}")> at {}'.format(self._uri,
                                                       hex(id(self)))

    @property
    @raise_if_connection_closed
    def conn(self):
        if self._conn is None:
            self._conn = _connect(**self._conn_kwargs)
            self._conn.set_session(
                psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED,
                autocommit=True)
            self._conn.set_client_encoding('UTF8')
            self._conn.autocommit = True
            cur = self._conn.cursor()
            cur.execute('SET TIME ZONE "UTC";')
            cur.execute('SET search_path TO %s;', (self.schema, ))
            cur.close()
        return self._conn

    @property
    def autocommit(self):
        return self.conn.autocommit

    @autocommit.setter
    def autocommit(self, value):
        self.conn.autocommit = value

    @property
    def is_closed(self):
        return self._closed

    @property
    def schema(self):
        return self._schema

    @raise_if_connection_closed
    def commit(self):
        self.conn.commit()

    @raise_if_connection_closed
    def rollback(self):
        self.conn.rollback()

    @raise_if_connection_closed
    def close(self):
        if self._conn is not None:
            self._conn.close()
            self._conn = None
        self._pg_pool = None
        self._closed = True

    @raise_if_connection_closed
    def cursor(self, named=False):
        if named:
            return self.conn.cursor(
                name=str(uuid.uuid4()), scrollable=True, withhold=True)
        return self.conn.cursor()

    @raise_if_connection_closed
    def reattach(self):
        if self._pg_pool is not None:
            self._pg_pool.unlock(self)
        elif self._autoclose:
            self.close()

    @raise_if_connection_closed
    def transaction(self, **kwargs):
        return PgTransaction(self, **kwargs)

class PgDbPool(object):

    __slots__ = ('_uri', '_pool_size', '_pool',
                 '_warmed_up', '_closed', '_lock')

    def __init__(self, uri, pool_size=5, **kwargs):
        self._uri = uri
        self._pool_size = pool_size
        self._pool = LockRing()
        self._warmed_up = False
        self._closed = False
        self._lock = threading.Lock()

    @property
    def is_closed(self):
        return self._closed

    @property
    def uri(self):
        return self._uri

    def __repr__(self):
        return '<PgDbPool("{}")> at {}'.format(self._uri, hex(id(self)))

    def __len__(self):
        return len(self._pool)

    def close(self):
        with self._lock:
            for __ in xrange(len(self._pool)):
                conn = self._pool.pop()
                conn.close()
            self._closed = True
            self._pool.reset()

    @raise_if_pool_closed
    def transaction(self, **kwargs):
        with self._lock:
            conn = self._get_conn()
        return conn.transaction(**kwargs)

    def _get_conn(self):
        if not self._warmed_up:
            self._warm_up()
        try:
            conn = self._pool.next()
        except (IndexError, StopIteration):
            logger.error('pool exhausted, making new connection')
            conn = self._new_conn(autoclose=True)
        return conn

    def _new_conn(self, **kwargs):
        return PgConnection(self._uri, **kwargs)

    @raise_if_pool_closed
    def unlock(self, conn):
        self._pool.push(conn)

    def _warm_up(self):
        self._warmed_up = True
        for __ in xrange(self._pool_size):
            self._pool.push(self._new_conn(pool=self))
        logger.info(
            'warmed up pool "%s" length %s', repr(self), self._pool_size)

def _connect(**kwargs):
    kwargs.setdefault('cursor_factory', psycopg2.extras.RealDictCursor)
    return psycopg2.connect(**kwargs)
<|endoftext|>