From 7c81940af582e4a3c8a2ce425c9b54a01f14c469 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Thu, 29 Dec 2022 14:50:55 +0100 Subject: [PATCH] ENH: add _load_requirements class attribute to Dataset classes --- yt/data_objects/static_output.py | 9 +++++++- yt/frontends/_skeleton/data_structures.py | 4 ++++ yt/frontends/arepo/data_structures.py | 1 + yt/frontends/cf_radial/data_structures.py | 1 + yt/frontends/chimera/data_structures.py | 1 + yt/frontends/cholla/data_structures.py | 1 + yt/frontends/chombo/data_structures.py | 5 +++-- yt/frontends/eagle/data_structures.py | 2 ++ yt/frontends/enzo/data_structures.py | 1 + yt/frontends/enzo_e/data_structures.py | 1 + yt/frontends/exodus_ii/data_structures.py | 1 + yt/frontends/fits/data_structures.py | 5 +++++ yt/frontends/flash/data_structures.py | 2 ++ yt/frontends/gadget/data_structures.py | 1 + yt/frontends/gadget_fof/data_structures.py | 1 + yt/frontends/gamer/data_structures.py | 1 + yt/frontends/gdf/data_structures.py | 1 + yt/frontends/gizmo/data_structures.py | 1 + yt/frontends/halo_catalog/data_structures.py | 1 + yt/frontends/http_stream/data_structures.py | 1 + yt/frontends/moab/data_structures.py | 1 + yt/frontends/nc4_cm1/data_structures.py | 1 + yt/frontends/open_pmd/data_structures.py | 2 ++ yt/frontends/owls/data_structures.py | 1 + yt/frontends/owls_subfind/data_structures.py | 1 + yt/frontends/ramses/data_structures.py | 1 + yt/frontends/sdf/data_structures.py | 1 + yt/frontends/swift/data_structures.py | 1 + yt/frontends/ytdata/data_structures.py | 9 ++++++++ yt/utilities/hierarchy_inspection.py | 22 +++++++++++++++++++- 30 files changed, 77 insertions(+), 4 deletions(-) diff --git a/yt/data_objects/static_output.py b/yt/data_objects/static_output.py index 9a7eed66cf..3ff2172040 100644 --- a/yt/data_objects/static_output.py +++ b/yt/data_objects/static_output.py @@ -142,7 +142,7 @@ def ireq(self, value): class Dataset(abc.ABC): - + _load_requirements: List[str] = [] default_fluid_type = "gas" default_field = ("gas", "density") fluid_types: Tuple[FieldType, ...] = ("gas", "deposit", "index") @@ -346,6 +346,13 @@ def force_periodicity(self, val=True): raise TypeError("force_periodicity expected a boolean.") self._force_periodicity = val + @classmethod + def _missing_extra_requirements(cls) -> List[str]: + # return a list of optional dependencies that are + # needed for the present class to function, but currently missing. + # returning an empty list means that all requirements are met + return [name for name in cls._load_requirements if not find_spec(name)] + # abstract methods require implementation in subclasses @classmethod @abc.abstractmethod diff --git a/yt/frontends/_skeleton/data_structures.py b/yt/frontends/_skeleton/data_structures.py index e3636bb75d..7c52fea323 100644 --- a/yt/frontends/_skeleton/data_structures.py +++ b/yt/frontends/_skeleton/data_structures.py @@ -1,5 +1,6 @@ import os import weakref +from typing import List import numpy as np @@ -76,6 +77,9 @@ class SkeletonDataset(Dataset): _index_class = SkeletonHierarchy _field_info_class = SkeletonFieldInfo + # names of additional modules that may be required to load data from disk + _load_requirements: List[str] = [] + def __init__( self, filename, diff --git a/yt/frontends/arepo/data_structures.py b/yt/frontends/arepo/data_structures.py index 042f0ab526..0b8dc621f4 100644 --- a/yt/frontends/arepo/data_structures.py +++ b/yt/frontends/arepo/data_structures.py @@ -8,6 +8,7 @@ class ArepoHDF5Dataset(GadgetHDF5Dataset): + _load_requirements = ["h5py"] _field_info_class = ArepoFieldInfo def __init__( diff --git a/yt/frontends/cf_radial/data_structures.py b/yt/frontends/cf_radial/data_structures.py index c76bb93e86..344621e23a 100644 --- a/yt/frontends/cf_radial/data_structures.py +++ b/yt/frontends/cf_radial/data_structures.py @@ -80,6 +80,7 @@ def _populate_grid_objects(self): class CFRadialDataset(Dataset): + _load_requirements = ["xarray", "pyart"] _index_class = CFRadialHierarchy _field_info_class = CFRadialFieldInfo diff --git a/yt/frontends/chimera/data_structures.py b/yt/frontends/chimera/data_structures.py index 08c86f564e..7dae9b5585 100644 --- a/yt/frontends/chimera/data_structures.py +++ b/yt/frontends/chimera/data_structures.py @@ -195,6 +195,7 @@ def _setup_data_io(self): class ChimeraDataset(Dataset): + _load_requirements = ["h5py"] _index_class = ChimeraUNSIndex # ChimeraHierarchy _field_info_class = ChimeraFieldInfo diff --git a/yt/frontends/cholla/data_structures.py b/yt/frontends/cholla/data_structures.py index a5e5723b91..59d4198df5 100644 --- a/yt/frontends/cholla/data_structures.py +++ b/yt/frontends/cholla/data_structures.py @@ -61,6 +61,7 @@ def _populate_grid_objects(self): class ChollaDataset(Dataset): + _load_requirements = ["h5py"] _index_class = ChollaHierarchy _field_info_class = ChollaFieldInfo diff --git a/yt/frontends/chombo/data_structures.py b/yt/frontends/chombo/data_structures.py index 874f4d7c2a..299e44c5a3 100644 --- a/yt/frontends/chombo/data_structures.py +++ b/yt/frontends/chombo/data_structures.py @@ -237,6 +237,7 @@ def _reconstruct_parent_child(self): class ChomboDataset(Dataset): + _load_requirements = ["h5py"] _index_class = ChomboHierarchy _field_info_class: Type[FieldInfoContainer] = ChomboFieldInfo @@ -641,7 +642,7 @@ def _read_particles(self): class Orion2Dataset(ChomboDataset): - + _load_requirements = ["h5py"] _index_class = Orion2Hierarchy _field_info_class = Orion2FieldInfo @@ -759,7 +760,7 @@ def __init__(self, ds, dataset_type="chombo_hdf5"): class ChomboPICDataset(ChomboDataset): - + _load_requirements = ["h5py"] _index_class = ChomboPICHierarchy _field_info_class = ChomboPICFieldInfo3D diff --git a/yt/frontends/eagle/data_structures.py b/yt/frontends/eagle/data_structures.py index 92046e4e3f..363635d75a 100644 --- a/yt/frontends/eagle/data_structures.py +++ b/yt/frontends/eagle/data_structures.py @@ -12,6 +12,7 @@ class EagleDataset(GadgetHDF5Dataset): + _load_requirements = ["h5py"] _particle_mass_name = "Mass" _field_info_class: Type[FieldInfoContainer] = OWLSFieldInfo _time_readin_ = "Time" @@ -71,6 +72,7 @@ def _is_valid(cls, filename, *args, **kwargs): class EagleNetworkDataset(EagleDataset): + _load_requirements = ["h5py"] _particle_mass_name = "Mass" _field_info_class = EagleNetworkFieldInfo _time_readin = "Time" diff --git a/yt/frontends/enzo/data_structures.py b/yt/frontends/enzo/data_structures.py index b0a107ac65..35d7c925f1 100644 --- a/yt/frontends/enzo/data_structures.py +++ b/yt/frontends/enzo/data_structures.py @@ -668,6 +668,7 @@ class EnzoDataset(Dataset): Enzo-specific output, set at a fixed time. """ + _load_requirements = ["h5py", "libconf"] _index_class = EnzoHierarchy _field_info_class = EnzoFieldInfo diff --git a/yt/frontends/enzo_e/data_structures.py b/yt/frontends/enzo_e/data_structures.py index a4d3f69ffd..d1e63efba4 100644 --- a/yt/frontends/enzo_e/data_structures.py +++ b/yt/frontends/enzo_e/data_structures.py @@ -286,6 +286,7 @@ class EnzoEDataset(Dataset): Enzo-E-specific output, set at a fixed time. """ + _load_requirements = ["h5py", "libconf"] refine_by = 2 _index_class = EnzoEHierarchy _field_info_class = EnzoEFieldInfo diff --git a/yt/frontends/exodus_ii/data_structures.py b/yt/frontends/exodus_ii/data_structures.py index 2372a247e1..6f1486467e 100644 --- a/yt/frontends/exodus_ii/data_structures.py +++ b/yt/frontends/exodus_ii/data_structures.py @@ -46,6 +46,7 @@ def _detect_output_fields(self): class ExodusIIDataset(Dataset): + _load_requirements = ["netCDF4"] _index_class = ExodusIIUnstructuredIndex _field_info_class = ExodusIIFieldInfo diff --git a/yt/frontends/fits/data_structures.py b/yt/frontends/fits/data_structures.py index 553d1b8901..97c0241c57 100644 --- a/yt/frontends/fits/data_structures.py +++ b/yt/frontends/fits/data_structures.py @@ -315,6 +315,7 @@ def check_sky_coords(filename, ndim): class FITSDataset(Dataset): + _load_requirements = ["astropy"] _index_class = FITSHierarchy _field_info_class: Type[FieldInfoContainer] = FITSFieldInfo _dataset_type = "fits" @@ -573,6 +574,7 @@ def find_axes(axis_names, prefixes): class YTFITSDataset(FITSDataset): + _load_requirements = ["astropy"] _field_info_class = YTFITSFieldInfo def _parse_parameter_file(self): @@ -655,6 +657,7 @@ def _is_valid(cls, filename, *args, **kwargs): class SkyDataFITSDataset(FITSDataset): + _load_requirements = ["astropy"] _field_info_class = WCSFITSFieldInfo def _determine_wcs(self): @@ -741,6 +744,7 @@ def _domain_decomp(self): class SpectralCubeFITSDataset(SkyDataFITSDataset): + _load_requirements = ["astropy"] _index_class = SpectralCubeFITSHierarchy def __init__( @@ -860,6 +864,7 @@ def _parse_index(self): class EventsFITSDataset(SkyDataFITSDataset): + _load_requirements = ["astropy"] _index_class = EventsFITSHierarchy def __init__( diff --git a/yt/frontends/flash/data_structures.py b/yt/frontends/flash/data_structures.py index 2362610a66..037f622af2 100644 --- a/yt/frontends/flash/data_structures.py +++ b/yt/frontends/flash/data_structures.py @@ -164,6 +164,7 @@ def _populate_grid_objects(self): class FLASHDataset(Dataset): + _load_requirements = ["h5py"] _index_class: Type[Index] = FLASHHierarchy _field_info_class = FLASHFieldInfo _handle = None @@ -484,6 +485,7 @@ class FLASHParticleFile(ParticleFile): class FLASHParticleDataset(FLASHDataset): + _load_requirements = ["h5py"] _index_class = ParticleIndex filter_bbox = False _file_class = FLASHParticleFile diff --git a/yt/frontends/gadget/data_structures.py b/yt/frontends/gadget/data_structures.py index 782fadd5c0..9c864e2839 100644 --- a/yt/frontends/gadget/data_structures.py +++ b/yt/frontends/gadget/data_structures.py @@ -574,6 +574,7 @@ class GadgetHDF5File(ParticleFile): class GadgetHDF5Dataset(GadgetDataset): + _load_requirements = ["h5py"] _file_class = GadgetHDF5File _index_class = SPHParticleIndex _field_info_class: Type[FieldInfoContainer] = GadgetFieldInfo diff --git a/yt/frontends/gadget_fof/data_structures.py b/yt/frontends/gadget_fof/data_structures.py index 5549dad4bd..aa423714f7 100644 --- a/yt/frontends/gadget_fof/data_structures.py +++ b/yt/frontends/gadget_fof/data_structures.py @@ -137,6 +137,7 @@ def _read_particle_positions(self, ptype, f=None): class GadgetFOFDataset(ParticleDataset): + _load_requirements = ["h5py"] _index_class = GadgetFOFParticleIndex _file_class = GadgetFOFHDF5File _field_info_class = GadgetFOFFieldInfo diff --git a/yt/frontends/gamer/data_structures.py b/yt/frontends/gamer/data_structures.py index 8dff75cdf2..e28c2df4c2 100644 --- a/yt/frontends/gamer/data_structures.py +++ b/yt/frontends/gamer/data_structures.py @@ -201,6 +201,7 @@ def _validate_parent_children_relationship(self): class GAMERDataset(Dataset): + _load_requirements = ["h5py"] _index_class = GAMERHierarchy _field_info_class = GAMERFieldInfo _handle = None diff --git a/yt/frontends/gdf/data_structures.py b/yt/frontends/gdf/data_structures.py index 818f9c80b5..27c63a776c 100644 --- a/yt/frontends/gdf/data_structures.py +++ b/yt/frontends/gdf/data_structures.py @@ -148,6 +148,7 @@ def _get_grid_children(self, grid): class GDFDataset(Dataset): + _load_requirements = ["h5py"] _index_class = GDFHierarchy _field_info_class = GDFFieldInfo diff --git a/yt/frontends/gizmo/data_structures.py b/yt/frontends/gizmo/data_structures.py index 63af45c8db..e886e3cd9b 100644 --- a/yt/frontends/gizmo/data_structures.py +++ b/yt/frontends/gizmo/data_structures.py @@ -7,6 +7,7 @@ class GizmoDataset(GadgetHDF5Dataset): + _load_requirements = ["h5py"] _field_info_class = GizmoFieldInfo @classmethod diff --git a/yt/frontends/halo_catalog/data_structures.py b/yt/frontends/halo_catalog/data_structures.py index 29a7f6fe1e..8e0c3018b9 100644 --- a/yt/frontends/halo_catalog/data_structures.py +++ b/yt/frontends/halo_catalog/data_structures.py @@ -98,6 +98,7 @@ class YTHaloCatalogDataset(SavedDataset): in yt_astro_analysis. """ + _load_requirements = ["h5py"] _index_class = ParticleIndex _file_class = YTHaloCatalogFile _field_info_class = YTHaloCatalogFieldInfo diff --git a/yt/frontends/http_stream/data_structures.py b/yt/frontends/http_stream/data_structures.py index d412056b61..47443658a7 100644 --- a/yt/frontends/http_stream/data_structures.py +++ b/yt/frontends/http_stream/data_structures.py @@ -15,6 +15,7 @@ class HTTPParticleFile(ParticleFile): class HTTPStreamDataset(ParticleDataset): + _load_requirements = ["requests"] _index_class = ParticleIndex _file_class = HTTPParticleFile _field_info_class = SPHFieldInfo diff --git a/yt/frontends/moab/data_structures.py b/yt/frontends/moab/data_structures.py index 2c01615ac0..46d520bf35 100644 --- a/yt/frontends/moab/data_structures.py +++ b/yt/frontends/moab/data_structures.py @@ -47,6 +47,7 @@ def _count_grids(self): class MoabHex8Dataset(Dataset): + _load_requirements = ["h5py"] _index_class = MoabHex8Hierarchy _field_info_class = MoabFieldInfo periodicity = (False, False, False) diff --git a/yt/frontends/nc4_cm1/data_structures.py b/yt/frontends/nc4_cm1/data_structures.py index 9cb31c237d..6fc1fe4604 100644 --- a/yt/frontends/nc4_cm1/data_structures.py +++ b/yt/frontends/nc4_cm1/data_structures.py @@ -64,6 +64,7 @@ def _populate_grid_objects(self): class CM1Dataset(Dataset): + _load_requirements = ["netCDF4"] _index_class = CM1Hierarchy _field_info_class = CM1FieldInfo diff --git a/yt/frontends/open_pmd/data_structures.py b/yt/frontends/open_pmd/data_structures.py index 978a5cf9a1..498777cf8b 100644 --- a/yt/frontends/open_pmd/data_structures.py +++ b/yt/frontends/open_pmd/data_structures.py @@ -431,6 +431,7 @@ class OpenPMDDataset(Dataset): - particle and mesh positions are *absolute* with respect to the simulation origin. """ + _load_requirements = ["h5py"] _index_class = OpenPMDHierarchy _field_info_class = OpenPMDFieldInfo @@ -656,6 +657,7 @@ def _load(self, it, **kwargs): class OpenPMDGroupBasedDataset(Dataset): + _load_requirements = ["h5py"] _index_class = OpenPMDHierarchy _field_info_class = OpenPMDFieldInfo diff --git a/yt/frontends/owls/data_structures.py b/yt/frontends/owls/data_structures.py index 498edd265f..08bca6d885 100644 --- a/yt/frontends/owls/data_structures.py +++ b/yt/frontends/owls/data_structures.py @@ -9,6 +9,7 @@ class OWLSDataset(GadgetHDF5Dataset): + _load_requirements = ["h5py"] _particle_mass_name = "Mass" _field_info_class = OWLSFieldInfo _time_readin = "Time_GYR" diff --git a/yt/frontends/owls_subfind/data_structures.py b/yt/frontends/owls_subfind/data_structures.py index 55eee57442..4edc873a1f 100644 --- a/yt/frontends/owls_subfind/data_structures.py +++ b/yt/frontends/owls_subfind/data_structures.py @@ -79,6 +79,7 @@ def __init__(self, ds, io, filename, file_id, bounds): class OWLSSubfindDataset(ParticleDataset): + _load_requirements = ["h5py"] _index_class = OWLSSubfindParticleIndex _file_class = OWLSSubfindHDF5File _field_info_class = OWLSSubfindFieldInfo diff --git a/yt/frontends/ramses/data_structures.py b/yt/frontends/ramses/data_structures.py index 84babeedf8..ab4072bd84 100644 --- a/yt/frontends/ramses/data_structures.py +++ b/yt/frontends/ramses/data_structures.py @@ -697,6 +697,7 @@ def print_stats(self): class RAMSESDataset(Dataset): + _load_requirements = ["f90nml"] _index_class = RAMSESIndex _field_info_class = RAMSESFieldInfo gamma = 1.4 # This will get replaced on hydro_fn open diff --git a/yt/frontends/sdf/data_structures.py b/yt/frontends/sdf/data_structures.py index 989292d9aa..ec0bf24001 100644 --- a/yt/frontends/sdf/data_structures.py +++ b/yt/frontends/sdf/data_structures.py @@ -24,6 +24,7 @@ class SDFFile(ParticleFile): class SDFDataset(ParticleDataset): + _load_requirements = ["requests"] _index_class = ParticleIndex _file_class = SDFFile _field_info_class = SDFFieldInfo diff --git a/yt/frontends/swift/data_structures.py b/yt/frontends/swift/data_structures.py index 1860ecca56..e1577b0c05 100644 --- a/yt/frontends/swift/data_structures.py +++ b/yt/frontends/swift/data_structures.py @@ -13,6 +13,7 @@ class SwiftParticleFile(ParticleFile): class SwiftDataset(SPHDataset): + _load_requirements = ["h5py"] _index_class = SPHParticleIndex _field_info_class = SPHFieldInfo _file_class = SwiftParticleFile diff --git a/yt/frontends/ytdata/data_structures.py b/yt/frontends/ytdata/data_structures.py index 2874c8e9fe..af819442f1 100644 --- a/yt/frontends/ytdata/data_structures.py +++ b/yt/frontends/ytdata/data_structures.py @@ -234,6 +234,7 @@ def __init__(self, ds, io, filename, file_id, range): class YTDataContainerDataset(YTDataset): """Dataset for saved geometric data containers.""" + _load_requirements = ["h5py"] _index_class = ParticleIndex _file_class = YTDataHDF5File _field_info_class: Type[FieldInfoContainer] = YTDataContainerFieldInfo @@ -316,6 +317,8 @@ def _is_valid(cls, filename, *args, **kwargs): class YTDataLightRayDataset(YTDataContainerDataset): """Dataset for saved LightRay objects.""" + _load_requirements = ["h5py"] + def _parse_parameter_file(self): super()._parse_parameter_file() self._restore_light_ray_solution() @@ -357,6 +360,7 @@ def _is_valid(cls, filename, *args, **kwargs): class YTSpatialPlotDataset(YTDataContainerDataset): """Dataset for saved slices and projections.""" + _load_requirements = ["h5py"] _field_info_class = YTGridFieldInfo def __init__(self, *args, **kwargs): @@ -477,6 +481,7 @@ def _populate_grid_objects(self): class YTGridDataset(YTDataset): """Dataset for saved covering grids, arbitrary grids, and FRBs.""" + _load_requirements = ["h5py"] _index_class: Type[Index] = YTGridHierarchy _field_info_class = YTGridFieldInfo _dataset_type = "ytgridhdf5" @@ -709,6 +714,7 @@ def _read_fluid_fields(self, fields, dobj, chunk=None): class YTNonspatialDataset(YTGridDataset): """Dataset for general array data.""" + _load_requirements = ["h5py"] _index_class = YTNonspatialHierarchy _field_info_class = YTGridFieldInfo _dataset_type = "ytnonspatialhdf5" @@ -777,6 +783,7 @@ def _is_valid(cls, filename, *args, **kwargs): class YTProfileDataset(YTNonspatialDataset): """Dataset for saved profile objects.""" + _load_requirements = ["h5py"] fluid_types = ("data", "gas", "standard_deviation") def __init__(self, filename, unit_system="cgs"): @@ -926,6 +933,8 @@ def __getitem__(self, field): class YTClumpTreeDataset(YTNonspatialDataset): """Dataset for saved clump-finder data.""" + _load_requirements = ["h5py"] + def __init__(self, filename, unit_system="cgs"): super().__init__(filename, unit_system=unit_system) self._load_tree() diff --git a/yt/utilities/hierarchy_inspection.py b/yt/utilities/hierarchy_inspection.py index 015e3c5d7b..ba6b9b9d79 100644 --- a/yt/utilities/hierarchy_inspection.py +++ b/yt/utilities/hierarchy_inspection.py @@ -1,7 +1,10 @@ import inspect from collections import Counter from functools import reduce -from typing import List, Optional, Type +from typing import Dict, List, Optional, Type + +from yt.data_objects.static_output import Dataset +from yt.utilities.object_registries import output_type_registry def find_lowest_subclasses( @@ -48,3 +51,20 @@ def find_lowest_subclasses( if hint is not None: retv = [x for x in retv if hint.lower() in x.__name__.lower()] return retv + + +MISSING_REQUIREMENTS: Dict[Type[Dataset], List[str]] + + +def get_classes_with_missing_requirements() -> Dict[Type[Dataset], List[str]]: + # encapsulates access to MISSING_REQUIREMENTS global in order to + # - only make the computation once + # - delay said computation until needed so that the result is independent of import order + if "MISSING_REQUIREMENTS" not in globals(): + global MISSING_REQUIREMENTS + MISSING_REQUIREMENTS = { + cls: missing + for cls in sorted(output_type_registry.values(), key=lambda c: c.__name__) + if (missing := cls._missing_extra_requirements()) + } + return MISSING_REQUIREMENTS