<commit_before>"""Hooks which override deep learning interfaces with remote execution functionality."""

import torch
import inspect
import re
import json
import types
import functools
from . import workers
from . import utils


class BaseHook(object):
    r""" A abstract interface for deep learning framework hooks."""

    def __init__(self):
        ""
    def __enter__(self):
        ""
    def __exit__(self):
        ""

class TorchHook(BaseHook):
    r""" A Hook which Overrides Methods on PyTorch Variables & Tensors -
    **Currently compatible with PyTorch 0.3.1**

    The purpose of this class is to:

        * extend torch methods to allow for the moving of tensors
          and variables from one worker to another
        * override torch methods to execute commands on one worker
          that are called on tensors controlled by the local worker.

    This class is typically the first thing you will initialize when
    using PySyft with PyTorch because it is responsible for augmenting
    PyTorch with PySyft's added functionality (such as remote execution).

    :Parameters:

        * **local_worker (**:class:`.workers.BaseWorker` **, optional)**
          you can optionally provide a local worker as a parameter which
          TorchHook will assume to be the worker owned by the local machine.
          If you leave it empty, TorchClient will automatically initialize
          a :class:`.workers.VirtualWorker` under the assumption you're
          looking to do local experimentation/development.

        * **is_client (bool, optional)** whether or not the TorchHook is
          being initialized as an end-user client. This can impact whether
          or not variables are deleted when they fall out of scope. If you set
          this incorrectly on a end user client, Tensors and Variables will
          never be deleted. If you set this incorrectly on a remote machine
          (not a client), tensors will not get saved. It's really only
          important if you're not initializing the local worker yourself. (Default: True)

        * **verbose (bool, optional)** whether or not to print operations
          as they occur. (Defalt: True)

    :Example:

    >>> from syft.core.hooks import TorchHook
    >>> from syft.core.hooks import torch
    >>> hook = TorchHook()
    Hooking into Torch...
    Overloading Complete.
    >>> x = torch.FloatTensor([1,2,3,4,5])
    >>> x
     1
     2
     3
     4
     5
    [torch.FloatTensor of size 5]
    """

    def __init__(self, local_worker=None, is_client=True, verbose=True):
        super().__init__()

        self.local_worker = local_worker
        if(self.local_worker is None):

            # Every TorchHook instance should have a local worker which is responsible for
            # interfacing with other workers. The worker interface is what allows the Torch
            # specific code in TorchHook to be agnostic to the means by which workers communicate
            # (such as peer-to-peer, sockets, through local ports, or all within the same process)

            if(hasattr(torch, 'local_worker')):
                self.local_worker = torch.local_worker
                if(verbose):
                    print("Torch seems to already have a local_worker object... \
                          using that one instead...")
            else:
                self.local_worker = workers.VirtualWorker(hook=self, is_client_worker=is_client)

        torch.local_worker = self.local_worker

        # this is a list of all module functions in the torch module
        self.torch_funcs = dir(torch)

        # this is the list of torch tensor types that we will override for remote execution
        self.tensor_types = [torch.FloatTensor,
                             torch.DoubleTensor,
                             torch.HalfTensor,
                             torch.ByteTensor,
                             torch.CharTensor,
                             torch.ShortTensor,
                             torch.IntTensor,
                             torch.LongTensor]

        # this is the list of torch tensor VARIABLE types that we will override for remote execution
        # Variables are simply tensors that are differentiable (support gradients)
        # Parameters are Variables that are also weights in a neural model
        self.var_types = [torch.autograd.variable.Variable, torch.nn.Parameter]

        # a list of all classes in which we will override their methods for remote execution
        self.tensorvar_types = self.tensor_types + [torch.autograd.variable.Variable]
        self.tensorvar_types_strs = [x.__name__ for x in self.tensorvar_types]
        self.tensorvar_methods = list(
            set(
                [method
                    for tensorvar in self.tensorvar_types
                    for method in dir(tensorvar)]
                )
            )

        # Methods that caused infinite recursion during testing
        # TODO: May want to handle the ones in "exclude" manually at
        #       some point
        self.exclude = (['ndimension', 'nelement', 'size', 'numel',
                         'type', 'tolist', 'dim', '__iter__', 'select'])

        # This one wasn't in dir(Variable) -- probably a C++ thing
        self.var_exclude = ['__getattr__']

        # Torch functions we don't want to override
        self.torch_exclude = ['save', 'load', 'typename']

        # Safety checks for serializing and deserializing torch objects
        # Desperately needs stress testing before going out in the wild
        self.map_tensor_type = {
            'torch.FloatTensor': torch.FloatTensor,
            'torch.DoubleTensor': torch.DoubleTensor,
            'torch.HalfTensor': torch.HalfTensor,
            'torch.ByteTensor': torch.ByteTensor,
            'torch.CharTensor': torch.CharTensor,
            'torch.ShortTensor': torch.ShortTensor,
            'torch.IntTensor': torch.IntTensor,
            'torch.LongTensor': torch.LongTensor
        }
        self.map_var_type = {
            'torch.autograd.variable.Variable': torch.autograd.variable.Variable,
            'torch.nn.parameter.Parameter': torch.nn.parameter.Parameter
        }
        self.map_torch_type = dict(self.map_tensor_type, **self.map_var_type)

        if(not hasattr(torch, 'hooked')):
            if(verbose):
                print('Hooking into Torch...')
            self._hook_torch_module()
            for t_type in self.tensor_types:
                self._hook_tensor(t_type)
            self._hook_variable()
            torch.hooked = True
            if(verbose):
                print('Overloading complete.')
        else:
            if(verbose):
                print("WARNING: Torch seems to be already overloaded... skipping...")
    
    def __exit__(self):
        import torch


    def types_guard(self, torch_type_str):
        """types_guard(torch_type_str) -> torch.Tensor or torch.autograd.Variable

        This method converts strings into a type reference. This prevents
        deserialized JSON from being able to instantiate objects of arbitrary
        type which would be a security concern.
<commit_msg>enable context manager to dehook; plans<commit_after>"""Hooks which override deep learning interfaces with remote execution functionality."""

import torch
import inspect
import re
import json
import types
import functools
import importlib
from . import workers
from . import utils


class BaseHook(object):
    r""" A abstract interface for deep learning framework hooks."""

    def __init__(self):
        ""
    def __enter__(self):
        ""
    def __exit__(self):
        ""

class TorchHook(BaseHook):
    r""" A Hook which Overrides Methods on PyTorch Variables & Tensors -
    **Currently compatible with PyTorch 0.3.1**

    The purpose of this class is to:

        * extend torch methods to allow for the moving of tensors
          and variables from one worker to another
        * override torch methods to execute commands on one worker
          that are called on tensors controlled by the local worker.

    This class is typically the first thing you will initialize when
    using PySyft with PyTorch because it is responsible for augmenting
    PyTorch with PySyft's added functionality (such as remote execution).

    :Parameters:

        * **local_worker (**:class:`.workers.BaseWorker` **, optional)**
          you can optionally provide a local worker as a parameter which
          TorchHook will assume to be the worker owned by the local machine.
          If you leave it empty, TorchClient will automatically initialize
          a :class:`.workers.VirtualWorker` under the assumption you're
          looking to do local experimentation/development.

        * **is_client (bool, optional)** whether or not the TorchHook is
          being initialized as an end-user client. This can impact whether
          or not variables are deleted when they fall out of scope. If you set
          this incorrectly on a end user client, Tensors and Variables will
          never be deleted. If you set this incorrectly on a remote machine
          (not a client), tensors will not get saved. It's really only
          important if you're not initializing the local worker yourself. (Default: True)

        * **verbose (bool, optional)** whether or not to print operations
          as they occur. (Defalt: True)

    :Example:

    >>> from syft.core.hooks import TorchHook
    >>> from syft.core.hooks import torch
    >>> hook = TorchHook()
    Hooking into Torch...
    Overloading Complete.
    >>> x = torch.FloatTensor([1,2,3,4,5])
    >>> x
     1
     2
     3
     4
     5
    [torch.FloatTensor of size 5]
    """

    def __init__(self, local_worker=None, is_client=True, verbose=True):
        super().__init__()

        self.local_worker = local_worker
        if(self.local_worker is None):

            # Every TorchHook instance should have a local worker which is responsible for
            # interfacing with other workers. The worker interface is what allows the Torch
            # specific code in TorchHook to be agnostic to the means by which workers communicate
            # (such as peer-to-peer, sockets, through local ports, or all within the same process)

            if(hasattr(torch, 'local_worker')):
                self.local_worker = torch.local_worker
                if(verbose):
                    print("Torch seems to already have a local_worker object... \
                          using that one instead...")
            else:
                self.local_worker = workers.VirtualWorker(hook=self, is_client_worker=is_client)

        torch.local_worker = self.local_worker

        # this is a list of all module functions in the torch module
        self.torch_funcs = dir(torch)

        # this is the list of torch tensor types that we will override for remote execution
        self.tensor_types = [torch.FloatTensor,
                             torch.DoubleTensor,
                             torch.HalfTensor,
                             torch.ByteTensor,
                             torch.CharTensor,
                             torch.ShortTensor,
                             torch.IntTensor,
                             torch.LongTensor]

        # this is the list of torch tensor VARIABLE types that we will override for remote execution
        # Variables are simply tensors that are differentiable (support gradients)
        # Parameters are Variables that are also weights in a neural model
        self.var_types = [torch.autograd.variable.Variable, torch.nn.Parameter]

        # a list of all classes in which we will override their methods for remote execution
        self.tensorvar_types = self.tensor_types + [torch.autograd.variable.Variable]
        self.tensorvar_types_strs = [x.__name__ for x in self.tensorvar_types]
        self.tensorvar_methods = list(
            set(
                [method
                    for tensorvar in self.tensorvar_types
                    for method in dir(tensorvar)]
                )
            )

        # Methods that caused infinite recursion during testing
        # TODO: May want to handle the ones in "exclude" manually at
        #       some point
        self.exclude = (['ndimension', 'nelement', 'size', 'numel',
                         'type', 'tolist', 'dim', '__iter__', 'select'])

        # This one wasn't in dir(Variable) -- probably a C++ thing
        self.var_exclude = ['__getattr__']

        # Torch functions we don't want to override
        self.torch_exclude = ['save', 'load', 'typename']

        # Safety checks for serializing and deserializing torch objects
        # Desperately needs stress testing before going out in the wild
        self.map_tensor_type = {
            'torch.FloatTensor': torch.FloatTensor,
            'torch.DoubleTensor': torch.DoubleTensor,
            'torch.HalfTensor': torch.HalfTensor,
            'torch.ByteTensor': torch.ByteTensor,
            'torch.CharTensor': torch.CharTensor,
            'torch.ShortTensor': torch.ShortTensor,
            'torch.IntTensor': torch.IntTensor,
            'torch.LongTensor': torch.LongTensor
        }
        self.map_var_type = {
            'torch.autograd.variable.Variable': torch.autograd.variable.Variable,
            'torch.nn.parameter.Parameter': torch.nn.parameter.Parameter
        }
        self.map_torch_type = dict(self.map_tensor_type, **self.map_var_type)

        if(not hasattr(torch, 'hooked')):
            if(verbose):
                print('Hooking into Torch...')
            self._hook_torch_module()
            for t_type in self.tensor_types:
                self._hook_tensor(t_type)
            self._hook_variable()
            torch.hooked = True
            if(verbose):
                print('Overloading complete.')
        else:
            if(verbose):
                print("WARNING: Torch seems to be already overloaded... skipping...")

    def __enter__(self):
        pass


    def __exit__(self):
        importlib.reload(torch)


    def types_guard(self, torch_type_str):
        """types_guard(torch_type_str) -> torch.Tensor or torch.autograd.Variable

        This method converts strings into a type reference. This prevents
        deserialized JSON from being able to instantiate objects of arbitrary
        type which would be a security concern.

{'fraction_non_alphanumeric': 0.07372282984909317, 'fraction_numerical': 0.0019382528035442337, 'mean_word_length': 2.135880182331235, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5480158', 'n_tokens_mistral': 3671, 'n_tokens_neox': 3572, 'n_words': 1333}