Skip to content

Commit

Permalink
Use global session
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Aug 13, 2024
1 parent dbb5c33 commit a2142ff
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 11 deletions.
3 changes: 0 additions & 3 deletions pygmt/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
from pygmt.accessors import GMTDataArrayAccessor
from pygmt.figure import Figure, set_display
from pygmt.io import load_dataarray
from pygmt.session_management import begin as _begin
from pygmt.session_management import end as _end
from pygmt.src import (
binstats,
Expand Down Expand Up @@ -66,7 +65,5 @@
xyz2grd,
)

# Start our global modern mode session
_begin()
# Tell Python to run _end when shutting down
_atexit.register(_end)
10 changes: 10 additions & 0 deletions pygmt/_state.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
"""
Private dictionary to keep tracking of current PyGMT state.
The feature is only meant for internal use by PyGMT and is experimental!
"""

_STATE = {
"session_name": None,
"module_calls": None,
}
14 changes: 14 additions & 0 deletions pygmt/clib/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import contextlib
import ctypes as ctp
import os
import pathlib
import sys
import warnings
Expand All @@ -17,6 +18,7 @@
import pandas as pd
import xarray as xr
from packaging.version import Version
from pygmt._state import _STATE
from pygmt.clib.conversion import (
array_to_datetime,
as_c_contiguous,
Expand All @@ -33,6 +35,7 @@
data_kind,
tempfile_from_geojson,
tempfile_from_image,
unique_name,
)

FAMILIES = [
Expand Down Expand Up @@ -209,6 +212,10 @@ def __enter__(self):
Calls :meth:`pygmt.clib.Session.create`.
"""
# This is the first time a Session object is created.
if _STATE["session_name"] is None:
# Set GMT_SESSION_NAME to a customized, unique value.
_STATE["session_name"] = os.environ["GMT_SESSION_NAME"] = unique_name()
self.create("pygmt-session")
return self

Expand Down Expand Up @@ -623,6 +630,12 @@ def call_module(self, module: str, args: str | list[str]):
GMTCLibError
If the returned status code of the function is non-zero.
"""
if _STATE["module_calls"] is None:
from pygmt.session_management import begin

_STATE["module_calls"] = []
begin()

c_call_module = self.get_libgmt_func(
"GMT_Call_Module",
argtypes=[ctp.c_void_p, ctp.c_char_p, ctp.c_int, ctp.c_void_p],
Expand Down Expand Up @@ -651,6 +664,7 @@ def call_module(self, module: str, args: str | list[str]):
"'args' must be either a string or a list of strings."
)

_STATE["module_calls"].append(module)
status = c_call_module(self.session_pointer, module.encode(), mode, argv)
if status != 0:
raise GMTCLibError(
Expand Down
15 changes: 7 additions & 8 deletions pygmt/session_management.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,8 @@
Modern mode session management modules.
"""

import os
import sys

from pygmt._state import _STATE
from pygmt.clib import Session
from pygmt.helpers import unique_name


def begin():
Expand All @@ -17,10 +14,6 @@ def begin():
Only meant to be used once for creating the global session.
"""
# On Windows, need to set GMT_SESSION_NAME to a unique value
if sys.platform == "win32":
os.environ["GMT_SESSION_NAME"] = unique_name()

prefix = "pygmt-session"
with Session() as lib:
lib.call_module(module="begin", args=[prefix])
Expand All @@ -37,5 +30,11 @@ def end():
background, convert them to the desired format (specified in
``pygmt.begin``), and bring the figures to the working directory.
"""
if _STATE["session_name"] is not None:
return

with Session() as lib:
lib.call_module(module="end", args=[])

# Reset the sesion name to None
_STATE["session_name"] = None

Check warning on line 40 in pygmt/session_management.py

View check run for this annotation

Codecov / codecov/patch

pygmt/session_management.py#L40

Added line #L40 was not covered by tests

0 comments on commit a2142ff

Please sign in to comment.