From 3e07d99b1e2184e6361cc1c1989da473986484f8 Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Mon, 23 Oct 2023 15:24:23 +0100
Subject: [PATCH] Slight tidying

---
 pydatalab/pydatalab/apps/raman/blocks.py      | 41 +++++++++++--------
 .../src/components/datablocks/RamanBlock.vue  |  2 +-
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/pydatalab/pydatalab/apps/raman/blocks.py b/pydatalab/pydatalab/apps/raman/blocks.py
index 3bed2f396..5347049b3 100644
--- a/pydatalab/pydatalab/apps/raman/blocks.py
+++ b/pydatalab/pydatalab/apps/raman/blocks.py
@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 from typing import List, Tuple
 
 import bokeh
@@ -23,20 +24,16 @@ def plot_functions(self):
         return (self.generate_raman_plot,)
 
     @classmethod
-    def load(self, location: str, ext: str) -> Tuple[pd.DataFrame, List[str]]:
+    def load(self, location: str | Path) -> Tuple[pd.DataFrame, List[str]]:
         if not isinstance(location, str):
             location = str(location)
-        print(location)
+        ext = os.path.splitext(location)[-1].lower()
+
         if ext == ".txt":
             df = pd.read_csv(location, sep=r"\s+")
             df = df.rename(columns={"#Wave": "wavenumber", "#Intensity": "intensity"})
-        if ext == ".wdf":
-            if self.test_1D_2D(location) == "1D":
-                df = self.make_wdf_df(location)
-            elif self.test_1D_2D(location) == "2D":
-                raise RuntimeError("This is 2D Raman mapping data")
-            else:
-                raise RuntimeError("Unrecognised .wdf file")
+        elif ext == ".wdf":
+            df = self.make_wdf_df(location)
 
         df["sqrt(intensity)"] = np.sqrt(df["intensity"])
         df["log(intensity)"] = np.log10(df["intensity"])
@@ -90,18 +87,27 @@ def load(self, location: str, ext: str) -> Tuple[pd.DataFrame, List[str]]:
         return df, y_options
 
     @classmethod
-    def test_1D_2D(self, location):
+    def make_wdf_df(self, location: Path | str) -> pd.DataFrame:
+        """Read the .wdf file with RosettaSciIO and try to extract
+        1D Raman spectra.
+
+        Parameters:
+            location: The location of the file to read.
+
+        Returns:
+            A dataframe with the appropriate columns.
+
+        """
+
         raman_data = file_reader(location)
+
         if len(raman_data[0]["axes"]) == 1:
-            return "1D"
+            pass
         elif len(raman_data[0]["axes"]) == 3:
-            return "2D"
+            raise RuntimeError("This block does not support 2D Raman yet.")
         else:
-            raise ValueError("Data is not compatible 1D or 2D Raman data")
+            raise RuntimeError("Data is not compatible 1D or 2D Raman data.")
 
-    @classmethod
-    def make_wdf_df(self, location):
-        raman_data = file_reader(location)
         intensity = raman_data[0]["data"]
         wavenumber_size = raman_data[0]["axes"][0]["size"]
         wavenumber_offset = raman_data[0]["axes"][0]["offset"]
@@ -128,8 +134,7 @@ def generate_raman_plot(self):
                     self.accepted_file_extensions,
                     ext,
                 )
-            # if I don't pass ext here and try generate it in the function it doesn't recognise data attribute
-            pattern_dfs, y_options = self.load(file_info["location"], ext)
+            pattern_dfs, y_options = self.load(file_info["location"])
             pattern_dfs = [pattern_dfs]
 
         if pattern_dfs:
diff --git a/webapp/src/components/datablocks/RamanBlock.vue b/webapp/src/components/datablocks/RamanBlock.vue
index 80db1250d..95146266f 100644
--- a/webapp/src/components/datablocks/RamanBlock.vue
+++ b/webapp/src/components/datablocks/RamanBlock.vue
@@ -6,7 +6,7 @@ DataBlockBase as a prop, and save from within DataBlockBase  -->
       v-model="file_id"
       :item_id="item_id"
       :block_id="block_id"
-      :extensions="['.txt','.wdf']"
+      :extensions="['.txt', '.wdf']"
       updateBlockOnChange
     />