Refactor module imports and file structure

mermaid/graph.py

@@ -8,7 +8,7 @@
 """
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
 from mermaid.configuration import Config
 
@@ -29,20 +29,25 @@ class Graph:
     script: str
     config: Optional[Config] = None
 
-    def save(self, path: Optional[Path] = None) -> None:
+    def save(self, path: Optional[Union[Path, str]] = None) -> None:
         """Save the diagram to a file.
 
         Args:
-            path (Optional[Path]): The path to save the diagram. If not
+            path (Optional[Union[Path,str]]): The path to save the diagram. If not
                 provided, the diagram will be saved in the current directory
                 with the title as the filename.
         Raises:
             ValueError: If the file extension is not '.mmd' or '.mermaid'.
         """
-        file_path: Path = path if path else Path(f'./{self.title}.mmd')
-        if file_path.suffix not in ['.mmd', '.mermaid']:
+
+        if path is None:
+            path = Path(f'./{self.title}.mmd')
+        if isinstance(path, str):
+            path = Path(path)
+
+        if path.suffix not in ['.mmd', '.mermaid']:
             raise ValueError("File extension must be '.mmd' or '.mermaid'")
-        with open(file_path, 'w') as file:
+        with open(path, 'w') as file:
             file.write(self.script)
 
     def _build_script(self) -> None:

pyproject.toml

@@ -13,7 +13,6 @@ homepage = "https://mermaidpy.vercel.app"
 [tool.poetry.dependencies]
 python = "^3.9"
 requests = "^2.31.0"
-ipython = "^8.17.2"
 
 
 [tool.poetry.group.dev.dependencies]

tests/test_graph.py

@@ -23,6 +23,10 @@ def test_save_graph_with_path(self) -> None:
         self.graph_test.save(Path('./test-graph-file.mmd'))
         self.assertTrue(Path.exists(Path('./test-graph-file.mmd')))
 
+    def test_save_graph_with_path_str(self) -> None:
+        self.graph_test.save('./test-graph-str-path.mmd')
+        self.assertTrue(Path.exists(Path('./test-graph-str-path.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.mmd'))
@@ -40,4 +44,6 @@ def tearDown(self) -> None:
             os.remove('./test-graph-file.mmd')
         if os.path.exists('./test graph file.mmd'):
             os.remove('./test graph file.mmd')
+        if os.path.exists('./test-graph-str-path.mmd'):
+            os.remove('./test-graph-str-path.mmd')
         return super().tearDown()