support enums

Author: Jammy2211

autoarray/mask/abstract_mask.py

@@ -3,8 +3,7 @@
 from abc import ABC
 import logging
 import numpy as np
-from pathlib import Path
-from typing import Dict, Union
+from typing import Dict
 
 from autoconf.fitsable import output_to_fits
 

autoarray/mask/mask_1d.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 import logging
 import numpy as np
 from pathlib import Path
@@ -172,7 +173,11 @@ def header_dict(self) -> Dict:
         -------
         A dictionary containing the pixel scale of the mask, which can be output to a .fits file.
         """
+        class GridKeys(Enum):
+            PIXSCA = "PIXSCA"
+            ORIGIN = "ORIGIN"
+
         return {
-            "PIXSCA": self.pixel_scales[0],
-            "ORIGIN": self.origin[0],
+            GridKeys.PIXSCA: self.pixel_scales[0],
+            GridKeys.ORIGIN: self.origin[0],
         }

autoarray/mask/mask_2d.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import logging
+from enum import Enum
 import numpy as np
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, List, Tuple, Union
@@ -716,11 +717,18 @@ def header_dict(self) -> Dict:
         -------
         A dictionary containing the pixel scale of the mask, which can be output to a .fits file.
         """
+
+        class GridKeys(Enum):
+            PIXSCAY = "PIXSCAY"
+            PIXSCAX = "PIXSCAX"
+            ORIGINY = "ORIGINY"
+            ORIGINX = "ORIGINX"
+
         return {
-            "PIXSCAY": self.pixel_scales[0],
-            "PIXSCAX": self.pixel_scales[1],
-            "ORIGINY": self.origin[0],
-            "ORIGINX": self.origin[1],
+            GridKeys.PIXSCAY: self.pixel_scales[0],
+            GridKeys.PIXSCAX: self.pixel_scales[1],
+            GridKeys.ORIGINY: self.origin[0],
+            GridKeys.ORIGINX: self.origin[1],
         }
 
     @property