Update supported file format for script saving to .mmd or .mermaid #40

mermaid/graph/__init__.py

@@ -15,16 +15,14 @@ class Graph:
     examples:
     >>> from mermaid.graph import Graph
     >>> grph = Graph("title","...")
-    >>> # save the script in txt file for other uses
-    >>> graph.save() # save in file with name of title and extention .txt.
-    >>> graph.svae('./file-name.txt') # save in specific file.
     """
     title: str
     script: str
 
     def save(self, path: Optional[Path] = None) -> None:
-        file_path: Path = path if path else Path(f'./{self.title}.txt')
-
+        file_path: Path = path if path else Path(f'./{self.title}.mmd')
+        if file_path.suffix not in ['.mmd', '.mermaid']:
+            raise ValueError("File extension must be '.mmd' or '.mermaid'")
         with open(file_path, 'w') as file:
             file.write(self.script)
 

mermaid/tests/test_graph.py

@@ -17,22 +17,27 @@ def setUp(self) -> None:
 
     def test_save_graph_without_path(self) -> None:
         self.graph_test.save()
-        self.assertTrue(Path.exists(Path('./test-graph.txt')))
+        self.assertTrue(Path.exists(Path('./test-graph.mmd')))
 
     def test_save_graph_with_path(self) -> None:
-        self.graph_test.save(Path('./test-graph-file.txt'))
-        self.assertTrue(Path.exists(Path('./test-graph-file.txt')))
+        self.graph_test.save(Path('./test-graph-file.mmd'))
+        self.assertTrue(Path.exists(Path('./test-graph-file.mmd')))
 
     def test_save_graph_name_contain_space(self) -> None:
         self.graph_test.title = 'test graph file'
-        self.graph_test.save(Path('./test graph file.txt'))
-        self.assertTrue(Path.exists(Path('./test graph file.txt')))
+        self.graph_test.save(Path('./test graph file.mmd'))
+        self.assertTrue(Path.exists(Path('./test graph file.mmd')))
+
+    def test_save_should_raise_valueerror(self):
+        self.graph_test.title = 'file-name'
+        with self.assertRaises(ValueError):
+            self.graph_test.save(Path('./file-name.txt'))
 
     def tearDown(self) -> None:
-        if os.path.exists('./test-graph.txt'):
-            os.remove('./test-graph.txt')
-        if os.path.exists('./test-graph-file.txt'):
-            os.remove('./test-graph-file.txt')
-        if os.path.exists('./test graph file.txt'):
-            os.remove('./test graph file.txt')
+        if os.path.exists('./test-graph.mmd'):
+            os.remove('./test-graph.mmd')
+        if os.path.exists('./test-graph-file.mmd'):
+            os.remove('./test-graph-file.mmd')
+        if os.path.exists('./test graph file.mmd'):
+            os.remove('./test graph file.mmd')
         return super().tearDown()

mermaid/tests/test_utils.py

@@ -7,7 +7,7 @@
 
 class TestUtils(unittest.TestCase):
     def setUp(self) -> None:
-        self.file_test: str = './test-graph.txt'
+        self.file_test: str = './test-graph.mmd'
         with open(self.file_test, 'w') as file:
             file.write("""graph TD;
     A-->B;
@@ -18,7 +18,7 @@ def setUp(self) -> None:
 
     def test_read_file_should_raise_error(self):
         with self.assertRaises(FileNotFoundError):
-            md.load('./file.txt')
+            md.load('./file.mmd')
 
     def test_read_file(self):
         graph: Graph = md.load(self.file_test)

test-graph-file.txt

@@ -0,0 +1,5 @@
+graph TD;
+    A-->B;
+    A-->C;
+    B-->D;
+    C-->D;