From bf6c1086cc909de2eec429d118a503e3f9d6fdcc Mon Sep 17 00:00:00 2001
From: Shyue Ping Ong <shyuep@gmail.com>
Date: Tue, 28 Jan 2025 15:41:49 -0800
Subject: [PATCH] Return dict output.

---
 src/pymatgen/core/structure.py | 35 +++++++++++++++++++++++++++-------
 tests/core/test_structure.py   | 13 +++++++++++--
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/src/pymatgen/core/structure.py b/src/pymatgen/core/structure.py
index 95e341cb1fe..af6a5e3c3d7 100644
--- a/src/pymatgen/core/structure.py
+++ b/src/pymatgen/core/structure.py
@@ -3352,8 +3352,8 @@ def get_symmetry_dataset(self, backend: Literal["moyopy"], **kwargs) -> moyopy.M
     def get_symmetry_dataset(self, backend: Literal["spglib"], **kwargs) -> spglib.SpglibDataset: ...
 
     def get_symmetry_dataset(
-        self, backend: Literal["moyopy", "spglib"] = "spglib", **kwargs
-    ) -> moyopy.MoyoDataset | spglib.SpglibDataset:
+        self, backend: Literal["moyopy", "spglib"] = "spglib", return_raw_dataset=False, **kwargs
+    ) -> dict | moyopy.MoyoDataset | spglib.SpglibDataset:
         """Get a symmetry dataset from the structure using either moyopy or spglib backend.
 
         If using the spglib backend (default), please cite:
@@ -3365,6 +3365,9 @@ def get_symmetry_dataset(
         Args:
             backend ("moyopy" | "spglib"): Which symmetry analysis backend to use.
                 Defaults to "spglib".
+            return_raw_dataset (bool): Whether to return the raw Dataset object from the backend. The default is
+                False, which returns a dict with a common subset of the data present in both datasets. If you use the
+                raw Dataset object, we do not guarantee that the format of the output is not going to change.
             **kwargs: Additional arguments passed to the respective backend's constructor.
                 For spglib, these are passed to SpacegroupAnalyzer (e.g. symprec, angle_tolerance).
                 For moyopy, these are passed to MoyoDataset constructor.
@@ -3376,6 +3379,9 @@ def get_symmetry_dataset(
             ImportError: If the requested backend is not installed.
             ValueError: If an invalid backend is specified.
         """
+        if backend not in ("moyopy", "spglib"):
+            raise ValueError(f"Invalid {backend=}, must be one of moyopy or spglib.")
+
         if backend == "moyopy":
             try:
                 import moyopy
@@ -3385,16 +3391,31 @@ def get_symmetry_dataset(
 
             # Convert structure to MoyoDataset format
             moyo_cell = moyopy.interface.MoyoAdapter.from_structure(self)
-            return moyopy.MoyoDataset(cell=moyo_cell, **kwargs)
+            dataset = moyopy.MoyoDataset(cell=moyo_cell, **kwargs)
 
-        if backend == "spglib":
+        else:
             from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
             sga = SpacegroupAnalyzer(self, **kwargs)
-            return sga.get_symmetry_dataset()
+            dataset = sga.get_symmetry_dataset()
+
+        if return_raw_dataset:
+            return dataset
+
+        dictdata = {k: getattr(dataset, k) for k in ("hall_number", "number", "site_symmetry_symbols", "wyckoffs")}
+
+        if backend == "spglib":
+            dictdata["international"] = dataset.international
+            dictdata["orbits"] = dataset.crystallographic_orbits
+            dictdata["std_origin_shift"] = dataset.origin_shift
+        else:
+            from pymatgen.symmetry.groups import SpaceGroup
+
+            dictdata["international"] = SpaceGroup.from_int_number(dataset.number).symbol
+            dictdata["orbits"] = dataset.orbits
+            dictdata["std_origin_shift"] = dataset.std_origin_shift
 
-        valid_backends = ("moyopy", "spglib")
-        raise ValueError(f"Invalid {backend=}, must be one of {valid_backends}")
+        return dictdata
 
 
 class IMolecule(SiteCollection, MSONable):
diff --git a/tests/core/test_structure.py b/tests/core/test_structure.py
index d8bb9f365f1..e5ececa44e6 100644
--- a/tests/core/test_structure.py
+++ b/tests/core/test_structure.py
@@ -969,7 +969,16 @@ def test_sites_setter(self):
     def test_get_symmetry_dataset(self):
         """Test getting symmetry dataset from structure using different backends."""
         # Test spglib backend
-        dataset = self.struct.get_symmetry_dataset(backend="spglib")
+        for backend in ("spglib", "moyopy"):
+            dataset = self.struct.get_symmetry_dataset(backend=backend)
+            assert isinstance(dataset, dict)
+            assert dataset["number"] == 227
+            assert dataset["international"] == "Fd-3m"
+            assert len(dataset["orbits"]) == 2
+            pytest.approx(dataset["std_origin_shift"], [0, 0, 0])
+
+        dataset = self.struct.get_symmetry_dataset(backend="spglib", return_raw_dataset=True)
+
         assert dataset.number == 227  # Fd-3m space group
         assert dataset.international == "Fd-3m"
         assert len(dataset.rotations) > 0
@@ -977,7 +986,7 @@ def test_get_symmetry_dataset(self):
 
         # Test moyopy backend if available
         moyopy = pytest.importorskip("moyopy")
-        dataset = self.struct.get_symmetry_dataset(backend="moyopy")
+        dataset = self.struct.get_symmetry_dataset(backend="moyopy", return_raw_dataset=True)
         assert isinstance(dataset, moyopy.MoyoDataset)
         assert dataset.prim_std_cell.numbers == [14, 14]  # Si atomic number is 14