optional suffix (#564)

* optional suffix

* docstring

* suffix defaults to .json

* replaces .json with suffix

* linters

Author: mekhlakapoor

src/aind_data_schema/base.py

@@ -129,7 +129,7 @@ def default_filename(cls):
         name = cls._get_direct_subclass(cls).__name__
         return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower() + ".json"
 
-    def write_standard_file(self, output_directory: Optional[Path] = None, prefix=None):
+    def write_standard_file(self, output_directory: Optional[Path] = None, prefix=None, suffix=None):
         """
         Writes schema to standard json file
         Parameters
@@ -138,13 +138,17 @@ def write_standard_file(self, output_directory: Optional[Path] = None, prefix=No
             optional Path object for output directory
         prefix:
             optional str for intended filepath with extra naming convention
-
+        suffix:
+            optional str for intended filepath with extra naming convention
         """
         if prefix is None:
             filename = self.default_filename()
         else:
             filename = str(prefix) + "_" + self.default_filename()
 
+        if suffix:
+            filename = filename.replace(".json", suffix)
+
         if output_directory is not None:
             output_directory = Path(output_directory)
             filename = output_directory / filename

tests/test_base.py

@@ -51,8 +51,8 @@ def test_default_filename(self, mock_log):
         )
 
     @patch("builtins.open", new_callable=unittest.mock.mock_open())
-    def test_write_standard_file_no_prefix(self, mock_open):
-        """tests that standard file is named and written as expected with no prefix"""
+    def test_write_standard_file_no_prefix_suffix(self, mock_open):
+        """tests that standard file is named and written as expected with no prefix or suffix"""
         p = Procedures.construct()
         default_filename = p.default_filename()
         json_contents = p.json(indent=3)
@@ -77,6 +77,21 @@ def test_write_standard_file_with_prefix(self, mock_open):
         mock_open.assert_called_once_with(expected_file_path, "w")
         mock_open.return_value.__enter__().write.assert_called_once_with(json_contents)
 
+    @patch("builtins.open", new_callable=unittest.mock.mock_open())
+    def test_write_standard_file_with_suffix(self, mock_open):
+        """tests that standard file is named and written as expected with filename suffix"""
+        p = Procedures.construct()
+        json_contents = p.json(indent=3)
+        suffix = ".aind.json"
+        p.write_standard_file(suffix=suffix)
+
+        # It's expected that the file will be written to something like
+        # procedure.aind.json
+        expected_file_path = "procedures.aind.json"
+
+        mock_open.assert_called_once_with(expected_file_path, "w")
+        mock_open.return_value.__enter__().write.assert_called_once_with(json_contents)
+
     @patch("builtins.open", new_callable=unittest.mock.mock_open())
     def test_write_standard_file_with_output_directory(self, mock_open):
         """tests that standard file is named and written as expected with designated output directory"""