<commit_before># Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

import os
import json

from notebook.base.handlers import APIHandler, IPythonHandler, json_errors
from notebook.utils import url_path_join

from tornado import gen, web
from tornado.concurrent import Future
from tornado.ioloop import IOLoop
from tornado.websocket import WebSocketHandler, websocket_connect
from tornado.httpclient import HTTPRequest
from tornado.escape import url_escape

from ipython_genutils.py3compat import cast_unicode
from jupyter_client.session import Session
from traitlets.config.configurable import LoggingConfigurable

# TODO: Find a better way to specify global configuration options 
# for a server extension.
KG_URL = os.getenv('KG_URL', 'http://127.0.0.1:8888/')
KG_HEADERS = json.loads(os.getenv('KG_HEADERS', '{}'))
KG_HEADERS.update({
    'Authorization': 'token {}'.format(os.getenv('KG_AUTH_TOKEN', ''))
})
VALIDATE_KG_CERT = os.getenv('VALIDATE_KG_CERT') not in ['no', 'false']

KG_CLIENT_KEY = os.getenv('KG_CLIENT_KEY')
KG_CLIENT_CERT = os.getenv('KG_CLIENT_CERT')
KG_CLIENT_CA = os.getenv('KG_CLIENT_CA')

KG_HTTP_USER = os.getenv('KG_HTTP_USER', '')
KG_HTTP_PASS = os.getenv('KG_HTTP_PASS', '')

# Get env variables to handle timeout of request and connection
KG_CONNECT_TIMEOUT = float(os.getenv('KG_CONNECT_TIMEOUT', 20.0))
KG_REQUEST_TIMEOUT = float(os.getenv('KG_REQUEST_TIMEOUT', 20.0))


class WebSocketChannelsHandler(WebSocketHandler, IPythonHandler):

    def set_default_headers(self):
        """Undo the set_default_headers in IPythonHandler

        which doesn't make sense for websockets
        """
        pass

    def get_compression_options(self):
        # use deflate compress websocket
        return {}

    def authenticate(self):
        """Run before finishing the GET request

        Extend this method to add logic that should fire before
        the websocket finishes completing.
        """
        # authenticate the request before opening the websocket
        if self.get_current_user() is None:
            self.log.warning("Couldn't authenticate WebSocket connection")
            raise web.HTTPError(403)

        if self.get_argument('session_id', False):
            self.session.session = cast_unicode(self.get_argument('session_id'))
        else:
            self.log.warning("No session ID specified")

    def initialize(self):
        self.log.debug("Initializing websocket connection %s", self.request.path)
        self.session = Session(config=self.config)
        # TODO: make kernel client class configurable
        self.client = KernelGatewayWSClient()

    @gen.coroutine
    def get(self, kernel_id, *args, **kwargs):
        self.authenticate()
        self.kernel_id = cast_unicode(kernel_id, 'ascii')
        super(WebSocketChannelsHandler, self).get(kernel_id=kernel_id, *args, **kwargs)

    def open(self, kernel_id, *args, **kwargs):
        '''Handle web socket connection open to notebook server.'''
        # Delegate to web socket handler
        self.client.on_open(
            kernel_id=kernel_id,
            message_callback=self.write_message,
            compression_options=self.get_compression_options()
        )

    def on_message(self, message):
        '''Forward message to web socket handler.'''
        self.client.on_message(message)

    def write_message(self, message):
        '''Send message back to client.'''
        super(WebSocketChannelsHandler, self).write_message(message)

    def on_close(self):
        self.log.debug("Closing websocket connection %s", self.request.path)
        self.client.on_close()
        super(WebSocketChannelsHandler, self).on_close()

class KernelGatewayWSClient(LoggingConfigurable):
    '''Proxy web socket connection to a kernel gateway.'''

    def __init__(self):
        self.ws = None
        self.ws_future = Future()

    @gen.coroutine
    def _connect(self, kernel_id):
        ws_url = url_path_join(
            KG_URL.replace('http', 'ws'), 
            '/api/kernels', 
            url_escape(kernel_id),
            'channels'
        )
        self.log.info('Connecting to {}'.format(ws_url))
        parameters = {
          "headers" : KG_HEADERS,
          "validate_cert" : VALIDATE_KG_CERT,
          "auth_username" : KG_HTTP_USER, 
          "auth_password" : <PASSWORD>,
          "connect_timeout" : KG_CONNECT_TIMEOUT,
          "request_timeout" : KG_REQUEST_TIMEOUT
        }
        if KG_CLIENT_KEY:
            parameters["client_key"] = KG_CLIENT_KEY
            parameters["client_cert"] = KG_CLIENT_CERT
            parameters["ca_certs"] = KG_CLIENT_CA
        request = HTTPRequest(ws_url, **parameters)
        self.ws_future = websocket_connect(request)
        self.ws = yield self.ws_future
        # TODO: handle connection errors/timeout

    def _disconnect(self):
        if self.ws is not None:
            # Close connection
            self.ws.close()
        elif not self.ws_future.done():
            # Cancel pending connection
            self.ws_future.cancel()

    @gen.coroutine
    def _read_messages(self, callback):<commit_msg>Update handlers.py<commit_after># Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

import os
import json

from notebook.base.handlers import APIHandler, IPythonHandler, json_errors
from notebook.utils import url_path_join

from tornado import gen, web
from tornado.concurrent import Future
from tornado.ioloop import IOLoop
from tornado.websocket import WebSocketHandler, websocket_connect
from tornado.httpclient import HTTPRequest
from tornado.escape import url_escape

from ipython_genutils.py3compat import cast_unicode
from jupyter_client.session import Session
from traitlets.config.configurable import LoggingConfigurable

# TODO: Find a better way to specify global configuration options 
# for a server extension.
KG_URL = os.getenv('KG_URL', 'http://127.0.0.1:8888/')
KG_HEADERS = json.loads(os.getenv('KG_HEADERS', '{}'))
KG_HEADERS.update({
    'Authorization': 'token {}'.format(os.getenv('KG_AUTH_TOKEN', ''))
})
VALIDATE_KG_CERT = os.getenv('VALIDATE_KG_CERT') not in ['no', 'false']

KG_CLIENT_KEY = os.getenv('KG_CLIENT_KEY')
KG_CLIENT_CERT = os.getenv('KG_CLIENT_CERT')
KG_CLIENT_CA = os.getenv('KG_CLIENT_CA')

KG_HTTP_USER = os.getenv('KG_HTTP_USER', '')
KG_HTTP_PASS = os.getenv('KG_HTTP_PASS', '')

# Get env variables to handle timeout of request and connection
KG_CONNECT_TIMEOUT = float(os.getenv('KG_CONNECT_TIMEOUT', 20.0))
KG_REQUEST_TIMEOUT = float(os.getenv('KG_REQUEST_TIMEOUT', 20.0))


class WebSocketChannelsHandler(WebSocketHandler, IPythonHandler):

    def set_default_headers(self):
        """Undo the set_default_headers in IPythonHandler

        which doesn't make sense for websockets
        """
        pass

    def get_compression_options(self):
        # use deflate compress websocket
        return {}

    def authenticate(self):
        """Run before finishing the GET request

        Extend this method to add logic that should fire before
        the websocket finishes completing.
        """
        # authenticate the request before opening the websocket
        if self.get_current_user() is None:
            self.log.warning("Couldn't authenticate WebSocket connection")
            raise web.HTTPError(403)

        if self.get_argument('session_id', False):
            self.session.session = cast_unicode(self.get_argument('session_id'))
        else:
            self.log.warning("No session ID specified")

    def initialize(self):
        self.log.debug("Initializing websocket connection %s", self.request.path)
        self.session = Session(config=self.config)
        # TODO: make kernel client class configurable
        self.client = KernelGatewayWSClient()

    @gen.coroutine
    def get(self, kernel_id, *args, **kwargs):
        self.authenticate()
        self.kernel_id = cast_unicode(kernel_id, 'ascii')
        super(WebSocketChannelsHandler, self).get(kernel_id=kernel_id, *args, **kwargs)

    def open(self, kernel_id, *args, **kwargs):
        '''Handle web socket connection open to notebook server.'''
        # Delegate to web socket handler
        self.client.on_open(
            kernel_id=kernel_id,
            message_callback=self.write_message,
            compression_options=self.get_compression_options()
        )

    def on_message(self, message):
        '''Forward message to web socket handler.'''
        self.client.on_message(message)

    def write_message(self, message):
        '''Send message back to client.'''
        super(WebSocketChannelsHandler, self).write_message(message)

    def on_close(self):
        self.log.debug("Closing websocket connection %s", self.request.path)
        self.client.on_close()
        super(WebSocketChannelsHandler, self).on_close()

class KernelGatewayWSClient(LoggingConfigurable):
    '''Proxy web socket connection to a kernel gateway.'''

    def __init__(self):
        self.ws = None
        self.ws_future = Future()

    @gen.coroutine
    def _connect(self, kernel_id):
        ws_url = url_path_join(
            KG_URL.replace('http', 'ws'), 
            '/api/kernels', 
            url_escape(kernel_id),
            'channels'
        )
        self.log.info('Connecting to {}'.format(ws_url))
        parameters = {
          "headers" : KG_HEADERS,
          "validate_cert" : VALIDATE_KG_CERT,
          "auth_username" : KG_HTTP_USER, 
          "auth_password" : <PASSWORD>,
          "connect_timeout" : KG_CONNECT_TIMEOUT,
          "request_timeout" : KG_REQUEST_TIMEOUT
        }
        if KG_CLIENT_KEY:
            parameters["client_key"] = KG_CLIENT_KEY
            parameters["client_cert"] = KG_CLIENT_CERT
            if KG_CLIENT_CA:
                parameters["ca_certs"] = KG_CLIENT_CA
        request = HTTPRequest(ws_url, **parameters)
        self.ws_future = websocket_connect(request)
        self.ws = yield self.ws_future
        # TODO: handle connection errors/timeout

    def _disconnect(self):
        if self.ws is not None:
            # Close connection
            self.ws.close()
        elif not self.ws_future.done():
            # Cancel pending connection
            self.ws_future.cancel()

    @gen.coroutine
    def _read_messages(self, callback):
{'fraction_non_alphanumeric': 0.0982469659025236, 'fraction_numerical': 0.0038528221922558272, 'mean_word_length': 3.0894052776683734, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3064923', 'n_tokens_mistral': 3144, 'n_tokens_neox': 3046, 'n_words': 726}