linting

Author: VicTrqt

src/atomate2/abinit/flows/dfpt.py

@@ -225,8 +225,7 @@ def make(
 @dataclass
 class ShgFlowMaker(DfptFlowMaker):
     """
-    Maker to generate a DFPT flow to compute the static nonlinear optical
-            susceptibility tensor for the second-harmonic generation with abinit.
+    Maker to compute the static NLO SHG tensor.
 
     The classmethods allow to tailor the flow for specific configurations.
 

src/atomate2/abinit/jobs/anaddb.py

@@ -22,6 +22,8 @@
 )
 
 if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
+
     from atomate2.abinit.utils.history import JobHistory
 
 logger = logging.getLogger(__name__)
@@ -153,8 +155,7 @@ def get_response(
 
 @dataclass
 class AnaddbDfptDteMaker(AnaddbMaker):
-    """Maker to create a job to analyze a DDB file with the utility anaddb
-        to get info from DFPT calculations with the DTE.
+    """Maker to get info from DFPT calculations (with DTE) from a merged DDB file.
 
     Parameters
     ----------

src/atomate2/abinit/jobs/response.py

@@ -47,11 +47,12 @@
 
 @dataclass
 class ResponseMaker(BaseAbinitMaker):
-    """Maker to create a job with a Response Function ABINIT calculation.
-        The type of RF is defined by the self.calc_type attribute.
+    """Maker for a Response Function ABINIT calculation job.
 
     Parameters
     ----------
+    calc_type : str
+        The type of RF.
     name : str
         The job name.
     """
@@ -258,9 +259,7 @@ def generate_dte_perts(
 @job
 def run_rf(
     perturbations: list[dict],
-    rf_maker: BaseAbinitMaker = field(
-        default_factory=BaseAbinitMaker
-    ),  # TODO: change to generic ResponseMaker
+    rf_maker: ResponseMaker,
     prev_outputs: list[str] | None = None,
 ) -> Flow:
     """

src/atomate2/abinit/run.py

@@ -10,7 +10,6 @@
 
 from atomate2 import SETTINGS
 from atomate2.abinit.utils.common import (
-    ANADDB_INPUT_FILE_NAME,
     INPUT_FILE_NAME,
     LOG_FILE_NAME,
     MRGDDB_INPUT_FILE_NAME,
@@ -142,7 +141,7 @@ def run_anaddb(
         max_end_time = start_time + wall_time
 
     with (
-        open(ANADDB_INPUT_FILE_NAME) as stdin,
+        # open(ANADDB_INPUT_FILE_NAME) as stdin,
         open(LOG_FILE_NAME, "w") as stdout,
         open(STDERR_FILE_NAME, "w") as stderr,
     ):

src/atomate2/abinit/schemas/anaddb.py

@@ -4,7 +4,7 @@
 
 import logging
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 
 # from typing import Type, TypeVar, Union, Optional, List
@@ -13,6 +13,7 @@
 import abipy.core.abinit_units as abu
 import numpy as np
 from abipy.dfpt.anaddbnc import AnaddbNcFile
+from abipy.dfpt.ddb import AnaddbError
 from abipy.flowtk import events
 from abipy.flowtk.utils import File
 from emmet.core.structure import StructureMetadata
@@ -152,7 +153,9 @@ def from_abinit_files(
             output_doc = CalculationOutput.from_abinit_anaddb(abinit_anaddb)
 
             completed_at = str(
-                datetime.fromtimestamp(os.stat(abinit_anaddb_file).st_mtime)
+                datetime.fromtimestamp(
+                    os.stat(abinit_anaddb_file).st_mtime, tz=timezone.utc
+                )
             )
 
         report = None
@@ -166,7 +169,7 @@ def from_abinit_files(
                 # report.run_completed is False even when it completed...
                 has_anaddb_completed = TaskState.SUCCESS
 
-        except Exception as exc:
+        except (AnaddbError, Exception) as exc:
             msg = f"{cls} exception while parsing event_report:\n{exc}"
             logger.critical(msg)
 

src/atomate2/abinit/schemas/calculation.py

@@ -292,7 +292,7 @@ def from_abinit_files(
                 update_data = output_doc_gsr.model_dump(exclude_unset=True)
                 output_doc = output_doc.model_copy(update=update_data, deep=True)
         if not output_doc:
-            raise Exception(f"Neither {abinit_out_file} nor {abinit_gsr_file}\
+            raise FileNotFoundError(f"Neither {abinit_out_file} nor {abinit_gsr_file}\
                     exists. This means that there is no output, \
                     which is not normal.")
 

src/atomate2/abinit/schemas/mrgddb.py

@@ -4,13 +4,13 @@
 
 import logging
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 
 # from typing import Type, TypeVar, Union, Optional, List
 from typing import Any, Optional, Union
 
-from abipy.dfpt.ddb import DdbFile
+from abipy.dfpt.ddb import DdbError, DdbFile
 from abipy.flowtk import events
 from abipy.flowtk.utils import File
 from emmet.core.structure import StructureMetadata
@@ -139,7 +139,9 @@ def from_abinit_files(
             output_doc = CalculationOutput.from_abinit_outddb(abinit_outddb)
 
             completed_at = str(
-                datetime.fromtimestamp(os.stat(abinit_outddb_file).st_mtime)
+                datetime.fromtimestamp(
+                    os.stat(abinit_outddb_file).st_mtime, tz=timezone.utc
+                )
             )
 
         report = None
@@ -153,7 +155,7 @@ def from_abinit_files(
                 # report.run_completed is False even when it completed...
                 has_mrgddb_completed = TaskState.SUCCESS
 
-        except Exception as exc:
+        except (DdbError, Exception) as exc:
             msg = f"{cls} exception while parsing event_report:\n{exc}"
             logger.critical(msg)
 

src/atomate2/abinit/sets/anaddb.py

@@ -7,13 +7,11 @@
 import logging
 import os
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable
 
 from abipy.abio.inputs import AnaddbInput
 from abipy.flowtk.utils import Directory
 from monty.json import MontyEncoder, jsanitize
-from pymatgen.io.abinit.abiobjects import KSampling
 from pymatgen.io.core import InputSet
 
 from atomate2.abinit.files import out_to_in
@@ -28,8 +26,10 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
+    from pathlib import Path
 
     from pymatgen.core.structure import Structure
+    from pymatgen.io.abinit.abiobjects import KSampling
 
 
 logger = logging.getLogger(__name__)
@@ -341,13 +341,41 @@ def get_anaddb_input(
         return generated_input
 
 
-def anaddbinp_dfpt_dte(structure, anaddb_kwargs=None):
-    anaddb_input = AnaddbInput.dfpt(
-        structure=structure, dte=True, anaddb_kwargs=anaddb_kwargs
-    )
-    return anaddb_input
+def anaddbinp_dfpt_dte(
+    structure: Structure, anaddb_kwargs: None | dict = None
+) -> AnaddbInput:
+    """
+    Generate the AnaddbInput to retrieve information from the DTE.
+
+    Parameters
+    ----------
+    structure
+        A structure.
+    anaddb_kwargs
+        A dictionary with additional anaddb keywords to set.
+
+    Returns
+    -------
+        An AnaddbInput
+
+    """
+    return AnaddbInput.dfpt(structure=structure, dte=True, anaddb_kwargs=anaddb_kwargs)
 
 
 @dataclass
 class AnaddbDfptDteInputGenerator(AnaddbInputGenerator):
+    """
+    A class to generate the AnaddbInput to retrieve information from the DTE.
+
+    Parameters
+    ----------
+    factory
+        A callable to generate the AnaddbInput for DTE DFPT.
+
+    Returns
+    -------
+        An AnaddbInput
+
+    """
+
     factory: Callable = anaddbinp_dfpt_dte

src/atomate2/abinit/sets/base.py

@@ -67,7 +67,7 @@ class AbiBroadInputGenerator(InputGenerator):
     """
 
     calc_type: str
-    prev_outputs_deps: tuple
+    prev_outputs_deps: str | tuple | None
 
     def check_format_prev_dirs(
         self, prev_dirs: str | tuple | list | Path | None
@@ -134,7 +134,7 @@ def resolve_dep_exts(prev_dir: str, exts: list[str]) -> tuple:
                     raise RuntimeError("Should not occur.")
                 if files is not None:
                     inp_files = [
-                        (f.path, AbinitInputGenerator._get_in_file_name(f.path))
+                        (f.path, AbiBroadInputGenerator._get_in_file_name(f.path))
                         for f in files
                     ]
                     irdvars = irdvars_for_ext(ext)
@@ -148,7 +148,7 @@ def resolve_dep_exts(prev_dir: str, exts: list[str]) -> tuple:
                 if os.path.exists(out_den):
                     irdvars = irdvars_for_ext("DEN")
                     inp_files.append(
-                        (out_den, AbinitInputGenerator._get_in_file_name(out_den))
+                        (out_den, AbiBroadInputGenerator._get_in_file_name(out_den))
                     )
                     break
                 last_timden = prev_outdir.find_last_timden_file()
@@ -165,7 +165,7 @@ def resolve_dep_exts(prev_dir: str, exts: list[str]) -> tuple:
                 irdvars = irdvars_for_ext(ext)
                 if out_file:
                     inp_files.append(
-                        (out_file, AbinitInputGenerator._get_in_file_name(out_file))
+                        (out_file, AbiBroadInputGenerator._get_in_file_name(out_file))
                     )
                     break
         else:

src/atomate2/abinit/sets/mrgddb.py

@@ -7,7 +7,6 @@
 import os
 import time
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from abipy.abio.input_tags import DDE, DTE
@@ -24,6 +23,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
+    from pathlib import Path
 
 __all__ = ["MrgddbInputSet", "MrgddbInputGenerator"]