dialog to create fit functions

src/gui_qt/fit/function_creation_dialog.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import inspect
 import numbers
-import textwrap
+import pathlib
+import re
 from typing import Any
 
 import numpy as np
@@ -14,15 +15,18 @@ from gui_qt.lib.namespace import QNamespaceWidget
 __all__ = ['QUserFitCreator']
 
 validator = QtGui.QRegExpValidator(QtCore.QRegExp('[A-Za-z]\S*'))
+pattern = re.compile(r'def func\(.*\):', flags=re.MULTILINE)
 
 
 class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
-    classCreated = QtCore.pyqtSignal(object)
+    classCreated = QtCore.pyqtSignal()
 
-    def __init__(self, parent=None):
+    def __init__(self, filepath: str|pathlib.Path, parent=None):
         super().__init__(parent=parent)
         self.setupUi(self)
 
+        self.filepath = pathlib.Path(filepath)
+
         self.description_widget = DescWidget(self)
         self.args_widget = ArgWidget(self)
         self.kwargs_widget = KwargsWidget(self)
@@ -44,10 +48,23 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
 
         self.update_function()
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, filepath: str|pathlib.Path):
+        self.filepath = pathlib.Path(filepath)
+
         return self
 
     def update_function(self):
+        prev_text = self.plainTextEdit.toPlainText().split('\n')
+        func_body = ''
+        in_body = False
+        for line in prev_text:
+            if in_body:
+                func_body += line
+                continue
+
+            if pattern.search(line) is not None:
+                in_body = True
+
         try:
             var = self.args_widget.get_parameter()
             var += self.kwargs_widget.get_parameter()
@@ -67,7 +84,12 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
             k += self.kwargs_widget.get_strings()
 
             k += '\n    @staticmethod\n'
-            k += f"    def func(x, {', '.join(var)}):\n"
+            if var:
+                k += f"    def func(x, {', '.join(var)}):\n"
+            else:
+                k += f'    def func(x):\n'
+
+            k += func_body
 
             self.plainTextEdit.setPlainText(k)
         except Exception as e:
@@ -85,6 +107,7 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
         ns = self.namespace_widget.namespace.namespace
         func_value = ns[invalue][0]
         ret_func = ''
+        import_name = ''
 
         if func_value is None:
             ret_func = invalue
@@ -99,15 +122,28 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
         else:
             f_string = ns[invalue][-1]
             args = f_string[f_string.find('('):]
+
             if inspect.ismethod(func_value):
                 ret_func = func_value.__self__.__name__ + '.func'+args
+                import_name = func_value.__self__.__name__
+
             elif hasattr(func_value, '__qualname__'):
-                ret_func = func_value.__qualname__.split('.')[0]
-            self._imports.add((inspect.getmodule(func_value).__name__, ret_func))
+                import_name = func_value.__qualname__.split('.')[0]
+                ret_func = func_value.__qualname__ + args
+
+            self._imports.add((inspect.getmodule(func_value).__name__, import_name))
 
         self.plainTextEdit.insertPlainText(ret_func)
         self.update_function()
 
+    def accept(self):
+        # maybe add a check for correctness
+
+        with self.filepath.open('a') as f:
+            f.write('\n\n')
+            f.write(self.plainTextEdit.toPlainText())
+        self.classCreated.emit()
+        super().accept()
 
 class KwargsWidget(QtWidgets.QWidget):
     Changed = QtCore.pyqtSignal()
@@ -363,7 +399,7 @@ class ArgWidget(QtWidgets.QWidget):
 
         return var
 
-    def get_strings(self):
+    def get_strings(self) -> str:
         args = []
         bnds = []
         for row in range(self.table.rowCount()):
@@ -434,7 +470,7 @@ class DescWidget(QtWidgets.QWidget):
             raise ValueError('Class name is empty')
         stringi = f'class {self.klass_lineedit.text()}:\n' \
                   f'    name = {self.name_lineedit.text()!r}\n' \
-                  f'    group = {self.group_lineedit.text()!r}\n' \
+                  f'    type = {self.group_lineedit.text()!r}\n' \
                   f'    equation = {self.eq_lineedit.text()!r}\n'
 
         return stringi