Skip to content

Commit

Permalink
PERF: delay expensive imports (netcdf4, importlib.metadata, multiproc…
Browse files Browse the repository at this point in the history
…essing, tarfile, tomlib, tomli_w)
  • Loading branch information
neutrinoceros committed Jun 18, 2023
1 parent b754a80 commit 0c611fe
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 12 deletions.
15 changes: 9 additions & 6 deletions yt/loaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,11 @@
import atexit
import os
import sys
import tarfile
import time
import types
import warnings
from importlib.metadata import entry_points
from multiprocessing import Pipe, Process
from multiprocessing.connection import Connection
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from urllib.parse import urlsplit

import numpy as np
Expand Down Expand Up @@ -44,6 +40,8 @@
)
from yt.utilities.on_demand_imports import _pooch as pooch, _ratarmount as ratarmount

if TYPE_CHECKING:
from multiprocessing.connection import Connection
# --- Loaders for known data formats ---


Expand Down Expand Up @@ -90,6 +88,8 @@ def load(
yt.utilities.exceptions.YTAmbiguousDataType
If the data format matches more than one class of similar specialization levels.
"""
from importlib.metadata import entry_points

fn = os.path.expanduser(fn)

if any(wildcard in fn for wildcard in "[]?!*"):
Expand Down Expand Up @@ -1503,6 +1503,7 @@ def load_sample(
- Corresponding sample data live at https://yt-project.org/data
"""
import tarfile

if fn is None:
print(
Expand Down Expand Up @@ -1636,7 +1637,7 @@ def safe_extract(tar, path=".", members=None, *, numeric_owner=False):


def _mount_helper(
archive: str, mountPoint: str, ratarmount_kwa: Dict, conn: Connection
archive: str, mountPoint: str, ratarmount_kwa: Dict, conn: "Connection"
):
try:
fuseOperationsObject = ratarmount.TarMount(
Expand Down Expand Up @@ -1697,6 +1698,8 @@ def load_archive(
- This function requires ratarmount to be installed.
- This function does not work on Windows system.
"""
import tarfile
from multiprocessing import Pipe, Process

warnings.warn(
"The 'load_archive' function is still experimental and may be unstable.",
Expand Down
12 changes: 6 additions & 6 deletions yt/utilities/configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,10 @@
from pathlib import Path
from typing import Callable, List

import tomli_w
from more_itertools import always_iterable

from yt.utilities.configuration_tree import ConfigLeaf, ConfigNode

if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib

configuration_callbacks: List[Callable[["YTConfig"], None]] = []


Expand Down Expand Up @@ -89,6 +83,10 @@ def read(self, file_names):
if not os.path.exists(fname):
continue
metadata = {"source": f"file: {fname}"}
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
try:
with open(fname, "rb") as fh:
data = tomllib.load(fh)
Expand All @@ -104,6 +102,8 @@ def read(self, file_names):
return file_names_read

def write(self, file_handler):
import tomli_w

value = self.config_root.as_dict()
config_as_str = tomli_w.dumps(value)

Expand Down
2 changes: 2 additions & 0 deletions yt/utilities/on_demand_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,8 @@ def __init__(self):
# imported first, can get file lock errors on some systems (including travis-ci)
# so we need to do this before initializing h5py_imports()!
# similar to this issue https://github.com/pydata/xarray/issues/2560
if find_spec("h5py") is None or find_spec("netCDF4") is None:
return
try:
import netCDF4 # noqa F401
except ImportError:
Expand Down

0 comments on commit 0c611fe

Please sign in to comment.