From f2271108b2a301c56cbf208dec220bbee47fff38 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 27 May 2024 11:04:44 +0200
Subject: [PATCH] Feat: input files in calculation module (#159)

* Expose input files via py4vasp.calculation
* Add test for calculation module
---
 src/py4vasp/calculation/__init__.py | 15 ++++++++++++---
 src/py4vasp/calculation/_class.py   |  8 ++++----
 tests/calculation/test_module.py    | 24 ++++++++++++++++++++++++
 tests/calculation/test_repr.py      |  2 +-
 4 files changed, 41 insertions(+), 8 deletions(-)
 create mode 100644 tests/calculation/test_module.py

diff --git a/src/py4vasp/calculation/__init__.py b/src/py4vasp/calculation/__init__.py
index 6558ec35..2e6857ce 100644
--- a/src/py4vasp/calculation/__init__.py
+++ b/src/py4vasp/calculation/__init__.py
@@ -39,11 +39,13 @@ class provides a more flexible interface with which you can determine the source
 In the latter example, you can change the path from which the data is extracted.
 """
 import importlib
+import pathlib
 
-from py4vasp import exception
+from py4vasp import control, exception
 from py4vasp._util import convert
 
-__all__ = (
+_input_files = ("INCAR", "KPOINTS", "POSCAR")
+_quantities = (
     "band",
     "bandgap",
     "born_effective_charge",
@@ -77,13 +79,20 @@ class provides a more flexible interface with which you can determine the source
     "workfunction",
 )
 _private = ("dispersion",)
+__all__ = _quantities + _input_files
+
+
+path = pathlib.Path(".")
 
 
 def __getattr__(attr):
-    if attr in (__all__ + _private):
+    if attr in (_quantities + _private):
         module = importlib.import_module(f"py4vasp.calculation._{attr}")
         class_ = getattr(module, convert.to_camelcase(attr))
         return class_.from_path(".")
+    elif attr in (_input_files):
+        class_ = getattr(control, attr)
+        return class_(".")
     else:
         message = f"Could not find {attr} in the possible attributes, please check the spelling"
         raise exception.MissingAttribute(message)
diff --git a/src/py4vasp/calculation/_class.py b/src/py4vasp/calculation/_class.py
index b91765ef..c0e4e3dc 100644
--- a/src/py4vasp/calculation/_class.py
+++ b/src/py4vasp/calculation/_class.py
@@ -123,7 +123,7 @@ def POSCAR(self, poscar):
 
 
 def _add_all_refinement_classes(calc, add_single_class):
-    for name in calculation.__all__:
+    for name in calculation._quantities:
         calc = add_single_class(calc, name)
     return calc
 
@@ -155,7 +155,7 @@ def _add_to_documentation(calc, name):
 
 
 def _add_input_files(calc):
-    calc._INCAR = control.INCAR(calc.path())
-    calc._KPOINTS = control.KPOINTS(calc.path())
-    calc._POSCAR = control.POSCAR(calc.path())
+    for name in calculation._input_files:
+        file_ = getattr(control, name)(calc.path())
+        setattr(calc, f"_{name}", file_)
     return calc
diff --git a/tests/calculation/test_module.py b/tests/calculation/test_module.py
new file mode 100644
index 00000000..52065114
--- /dev/null
+++ b/tests/calculation/test_module.py
@@ -0,0 +1,24 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import pytest
+
+from py4vasp import Calculation, calculation, exception
+
+
+def test_access_of_attributes():
+    calc = Calculation.from_path(".")
+    for key in filter(attribute_included, dir(calc)):
+        getattr(calculation, key)
+
+
+def attribute_included(attr):
+    if attr.startswith("_"):  # do not include private attributes
+        return False
+    if attr.startswith("from"):  # do not include classmethods
+        return False
+    return True
+
+
+def test_nonexisting_attribute():
+    with pytest.raises(exception.MissingAttribute):
+        calculation.does_not_exist
diff --git a/tests/calculation/test_repr.py b/tests/calculation/test_repr.py
index 19e24d78..205ec3d0 100644
--- a/tests/calculation/test_repr.py
+++ b/tests/calculation/test_repr.py
@@ -7,7 +7,7 @@
 
 
 def test_repr():
-    for name in calculation.__all__:
+    for name in calculation._quantities:
         instance = getattr(calculation, name)
         class_name = convert.to_camelcase(name)
         module = importlib.import_module(f"py4vasp.calculation._{name}")