Skip to content

Commit

Permalink
avoid jupyter_client.session monkeypatch at import time
Browse files Browse the repository at this point in the history
patch in App, Client classes

this will hopefully shift to Session config in the future.
  • Loading branch information
minrk committed Mar 10, 2024
1 parent 9ff1800 commit 878f297
Show file tree
Hide file tree
Showing 9 changed files with 19 additions and 15 deletions.
3 changes: 3 additions & 0 deletions ipyparallel/apps/baseapp.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
from traitlets import Bool, Instance, Unicode, default, observe
from traitlets.config.application import LevelFormatter, catch_config_error

from ipyparallel import util

from .._version import __version__

# FIXME: CUnicode is needed for cli parsing
Expand Down Expand Up @@ -136,6 +138,7 @@ def _default_session(self):
@catch_config_error
def initialize(self, argv=None):
"""initialize the app"""
util._disable_session_extract_dates()
self.init_config_from_env()
super().initialize(argv)
self.init_deprecated_config()
Expand Down
1 change: 1 addition & 0 deletions ipyparallel/client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -570,6 +570,7 @@ def __init__(
)
raise ValueError(msg.format(exc.message))

util._disable_session_extract_dates()
self.session = Session(**extra_args)

self._query_socket = self._context.socket(zmq.DEALER)
Expand Down
2 changes: 2 additions & 0 deletions ipyparallel/controller/heartmonitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from zmq.devices import ThreadDevice, ThreadMonitoredQueue
from zmq.eventloop.zmqstream import ZMQStream

from ipyparallel import util
from ipyparallel.util import bind, connect, log_errors, set_hwm


Expand Down Expand Up @@ -120,6 +121,7 @@ class HeartMonitor(LoggingConfigurable):

@default("session")
def _default_session(self):
util._disable_session_extract_dates()
return Session(parent=self)

loop = Instance(ioloop.IOLoop)
Expand Down
1 change: 1 addition & 0 deletions ipyparallel/controller/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def _default_loop(self):

@default("session")
def _default_session(self):
util._disable_session_extract_dates()
return jupyter_client.session.Session(parent=self)

client_stream = Instance(
Expand Down
2 changes: 2 additions & 0 deletions ipyparallel/engine/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from traitlets.config import Config
from zmq.eventloop import zmqstream

from ipyparallel import util
from ipyparallel.apps.baseapp import (
BaseParallelApplication,
base_aliases,
Expand Down Expand Up @@ -405,6 +406,7 @@ def load_connection_file(self):

config.Session.packer = d['pack']
config.Session.unpacker = d['unpack']
util._disable_session_extract_dates()
self.session = Session(parent=self)

self.log.debug("Config changed:")
Expand Down
1 change: 1 addition & 0 deletions ipyparallel/engine/nanny.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ def __init__(
self.curve_secretkey = curve_secretkey
self.config = config
self.pipe = pipe
util._disable_session_extract_dates()
self.session = Session(config=self.config)

self.log = local_logger(f"{self.__class__.__name__}.{engine_id}", log_level)
Expand Down
1 change: 1 addition & 0 deletions ipyparallel/tests/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

class TaskDBTest:
def setUp(self):
util._disable_session_extract_dates()
self.session = Session()
self.db = self.create_db()
self.load_records(16)
Expand Down
21 changes: 7 additions & 14 deletions ipyparallel/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
from IPython import get_ipython
from IPython.core.profiledir import ProfileDir, ProfileDirError
from IPython.paths import get_ipython_dir
from jupyter_client import session
from jupyter_client.localinterfaces import is_public_ip, localhost, public_ips
from tornado.ioloop import IOLoop
from traitlets.log import get_logger
Expand Down Expand Up @@ -599,21 +600,13 @@ def _v(version_s):
return tuple(int(s) for s in re.findall(r"\d+", version_s))


def _patch_jupyter_client_dates():
"""Monkeypatch jupyter_client.extract_dates to be nondestructive wrt timezone info"""
import jupyter_client

if _v(jupyter_client.__version__) < _v('5.0'):
from jupyter_client import session

if hasattr(session, '_save_extract_dates'):
return
session._save_extract_dates = session.extract_dates
session.extract_dates = extract_dates

@lru_cache()
def _disable_session_extract_dates():
"""Monkeypatch jupyter_client.extract_dates to be a no-op
# FIXME: remove patch when we require jupyter_client 5.0
_patch_jupyter_client_dates()
avoids performance problem parsing unused timestamp strings
"""
session.extract_dates = lambda obj: obj


def progress(*args, widget=None, **kwargs):
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ dependencies = [
"pyzmq>=18",
"traitlets>=4.3",
"ipython>=4",
"jupyter_client",
"jupyter_client>=5",
"ipykernel>=4.4",
"tornado>=5.1",
"psutil",
Expand Down

0 comments on commit 878f297

Please sign in to comment.