diff --git a/src/gui_qt/_py/fitcreationdialog.py b/src/gui_qt/_py/fitcreationdialog.py
index d1c6075..4eeceec 100644
--- a/src/gui_qt/_py/fitcreationdialog.py
+++ b/src/gui_qt/_py/fitcreationdialog.py
@@ -2,7 +2,7 @@
 
 # Form implementation generated from reading ui file 'resources/_ui/fitcreationdialog.ui'
 #
-# Created by: PyQt5 UI code generator 5.15.4
+# Created by: PyQt5 UI code generator 5.15.10
 #
 # WARNING: Any manual changes made to this file will be lost when pyuic5 is
 # run again.  Do not edit this file unless you know what you are doing.
@@ -51,9 +51,8 @@ class Ui_Dialog(object):
         self.verticalLayout_6 = QtWidgets.QVBoxLayout(self.namespace_box)
         self.verticalLayout_6.setObjectName("verticalLayout_6")
         self.tabWidget.addTab(self.namespace_box, "")
-        self.plainTextEdit = CodeEditor(self.splitter)
-        self.plainTextEdit.setEnabled(True)
-        self.plainTextEdit.setObjectName("plainTextEdit")
+        self.editor = EditorWidget(self.splitter)
+        self.editor.setObjectName("editor")
         self.verticalLayout.addWidget(self.splitter)
         self.buttonBox = QtWidgets.QDialogButtonBox(Dialog)
         self.buttonBox.setOrientation(QtCore.Qt.Horizontal)
@@ -63,8 +62,8 @@ class Ui_Dialog(object):
 
         self.retranslateUi(Dialog)
         self.tabWidget.setCurrentIndex(0)
-        self.buttonBox.accepted.connect(Dialog.accept)
-        self.buttonBox.rejected.connect(Dialog.reject)
+        self.buttonBox.accepted.connect(Dialog.accept) # type: ignore
+        self.buttonBox.rejected.connect(Dialog.reject) # type: ignore
         QtCore.QMetaObject.connectSlotsByName(Dialog)
 
     def retranslateUi(self, Dialog):
@@ -74,4 +73,4 @@ class Ui_Dialog(object):
         self.tabWidget.setTabText(self.tabWidget.indexOf(self.args_box), _translate("Dialog", "Variables"))
         self.tabWidget.setTabText(self.tabWidget.indexOf(self.kwargs_box), _translate("Dialog", "Multiple choice"))
         self.tabWidget.setTabText(self.tabWidget.indexOf(self.namespace_box), _translate("Dialog", "Available Functions"))
-from ..lib.codeeditor import CodeEditor
+from ..lib.codeeditor import EditorWidget
diff --git a/src/gui_qt/fit/function_creation_dialog.py b/src/gui_qt/fit/function_creation_dialog.py
index 9dddd49..25562cd 100644
--- a/src/gui_qt/fit/function_creation_dialog.py
+++ b/src/gui_qt/fit/function_creation_dialog.py
@@ -54,7 +54,7 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
         return self
 
     def update_function(self):
-        prev_text = self.plainTextEdit.toPlainText().split('\n')
+        prev_text = self.editor.toPlainText().split('\n')
         func_body = ''
         in_body = False
         for line in prev_text:
@@ -89,9 +89,12 @@ class QUserFitCreator(QtWidgets.QDialog, Ui_Dialog):
             else:
                 k += f'    def func(x):\n'
 
-            k += func_body
+            if func_body:
+                k += func_body
+            else:
+                k += '        return x'
 
-            self.plainTextEdit.setPlainText(k)
+            self.editor.setPlainText(k)
         except Exception as e:
             QtWidgets.QMessageBox.warning(self, 'Failure', f'Error found: {e.args[0]}')
 
@@ -475,12 +478,3 @@ class DescWidget(QtWidgets.QWidget):
                   f"    equation = r'{self.eq_lineedit.text()}'\n"
 
         return stringi
-
-
-if __name__ == '__main__':
-    import sys
-    app = QtWidgets.QApplication([])
-    win = QUserFitCreator()
-    win.show()
-
-    sys.exit(app.exec())
diff --git a/src/gui_qt/lib/codeeditor.py b/src/gui_qt/lib/codeeditor.py
index 5bf0e04..f910b26 100644
--- a/src/gui_qt/lib/codeeditor.py
+++ b/src/gui_qt/lib/codeeditor.py
@@ -72,7 +72,8 @@ class PythonHighlighter(QtGui.QSyntaxHighlighter):
             (r'\bdef\b\s*(\w+)', 1, STYLES['defclass']),
             # 'class' followed by an identifier
             (r'\bclass\b\s*(\w+)', 1, STYLES['defclass']),
-            # @ followed by a word
+
+            # decorator @ followed by a word
             (r'\s*@(\w+)\s*', 0, STYLES['property']),
 
             # Numeric literals
@@ -80,7 +81,6 @@ class PythonHighlighter(QtGui.QSyntaxHighlighter):
             (r'\b[+-]?0[xX][\dA-Fa-f]+[lL]?\b', 0, STYLES['numbers']),
             (r'\b[+-]?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?\b', 0, STYLES['numbers']),
 
-
             # Double-quoted string, possibly containing escape sequences
             (r'[rf]?"[^"\\]*(\\.[^"\\]*)*"', 0, STYLES['string']),
             # Single-quoted string, possibly containing escape sequences
@@ -185,7 +185,6 @@ class CodeEditor(QtWidgets.QPlainTextEdit):
         self.update_width_linenumber(0)
 
         self.highlight = PythonHighlighter(self.document())
-        self.textChanged.connect(self._check_syntax)
 
     def keyPressEvent(self, evt):
         if evt.key() == QtCore.Qt.Key_Tab:
@@ -263,22 +262,48 @@ class CodeEditor(QtWidgets.QPlainTextEdit):
 
         self.setExtraSelections(extra_selections)
 
-    def color_line(self, color):
-        # is_valid, exception = self._check_syntax()
-        # if is_valid == 1:
-        doc = self.document()
-        print(doc.findBlockByLineNumber(color))
+
+class EditorWidget(QtWidgets.QWidget):
+    def __init__(self, parent=None):
+        super().__init__(parent=parent)
+
+        layout = QtWidgets.QVBoxLayout()
+        layout.setContentsMargins(0, 0, 0, 0)
+
+        self.editor = CodeEditor(self)
+        layout.addWidget(self.editor)
+
+        self.error_label = QtWidgets.QLabel(self)
+
+        font = QtGui.QFont()
+        font.setBold(True)
+        font.setWeight(75)
+        self.error_label.setFont(font)
+
+        self.error_label.setVisible(False)
+
+        layout.addWidget(self.error_label)
+
+        self.setLayout(layout)
+
+        for attr in ['appendPlainText', 'toPlainText', 'insertPlainText', 'setPlainText']:
+            setattr(self, attr, getattr(self.editor, attr))
+
+        self.editor.textChanged.connect(self._check_syntax)
 
     def _check_syntax(self) -> (int, tuple[typing.Any]):
+        is_valid = True
+
         # Compile into an AST and check for syntax errors.
         try:
             _ = parse(self.toPlainText(), filename='<string>')
-        except SyntaxError as e:
-            print('SyntaxError', e, e.args[0], e.lineno, e.offset, e.text)
-            self.color_line(e.lineno)
-            return 1, (e.lineno, e.offset)
-        except Exception as e:
-            print('Unexpected error', e)
-            return 2, (e.args[0],)
 
-        return 0, tuple()
+        except SyntaxError as e:
+            self.error_label.setText(f'Syntax error in line {e.lineno}: {e.args[0]}')
+            is_valid = False
+
+        except Exception as e:
+            self.error_label.setText(f'Unexpected error: {e.args[0]}')
+            is_valid = False
+
+        self.error_label.setVisible(not is_valid)
diff --git a/src/resources/_ui/fitcreationdialog.ui b/src/resources/_ui/fitcreationdialog.ui
index f6893a3..f2b9f03 100644
--- a/src/resources/_ui/fitcreationdialog.ui
+++ b/src/resources/_ui/fitcreationdialog.ui
@@ -86,11 +86,7 @@
        <layout class="QVBoxLayout" name="verticalLayout_6"/>
       </widget>
      </widget>
-     <widget class="CodeEditor" name="plainTextEdit">
-      <property name="enabled">
-       <bool>true</bool>
-      </property>
-     </widget>
+     <widget class="EditorWidget" name="editor" native="true"/>
     </widget>
    </item>
    <item>
@@ -107,9 +103,10 @@
  </widget>
  <customwidgets>
   <customwidget>
-   <class>CodeEditor</class>
-   <extends>QPlainTextEdit</extends>
+   <class>EditorWidget</class>
+   <extends>QWidget</extends>
    <header>..lib.codeeditor</header>
+   <container>1</container>
   </customwidget>
  </customwidgets>
  <resources/>