From d409e3b73317c0e3dee058347920c6760e1964d1 Mon Sep 17 00:00:00 2001
From: Paulo Henrique Garcia <paulohenriquesg@gmail.com>
Date: Thu, 2 Jun 2022 10:04:32 +0200
Subject: [PATCH] Add return type and method overwrite logic

---
 README.md     | 10 ++++++++
 src/class.js  | 13 +++++++---
 src/method.js | 70 ++++++++++++++++++++++++++++++++++++++++-----------
 3 files changed, 75 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index ddfd8b1..68d6fae 100644
--- a/README.md
+++ b/README.md
@@ -74,3 +74,13 @@ myPhpFile.findFunction('foo').setCode('return $a + $b');
 var myPhpFile = new writer(contents, options);
 myPhpFile.findClass('Foo_Class').setProperty('bar', 'foo-bar');
 ```
+
+##### Find a class and create a new method:
+
+If the method is not present on the class, it will be created. Otherwise, it will be overwritten.
+
+```js
+var myPhpFile = new writer(contents, options);
+myPhpFile.findClass('Foo_Class').setMethod('test', 'int $test = 0', 'return $test;', 'private', '?int');
+```
+
diff --git a/src/class.js b/src/class.js
index 889d6db..be7205c 100644
--- a/src/class.js
+++ b/src/class.js
@@ -156,12 +156,18 @@ Class.prototype.getMethod = function(name) {
 /**
  * Appends or update an function
  */
-Class.prototype.setMethod = function(name, args, body, flags) {
+Class.prototype.setMethod = function(name, args, body, flags, returnType) {
+    if (typeof returnType !== 'undefined') {
+        if (returnType[0] !== ':') {
+            returnType = ':' + returnType;
+        }
+    }
+    
     var method = this.getMethod(name);
     if (!method) {
         // append the function
         var ast = parser.parseEval('class a { \n' +
-            flags + ' function ' + name + '(' + args +  ') {\n' +
+            flags + ' function ' + name + '(' + args +  ') '+ returnType + '{\n' +
                 body + '\n' +
             '}\n' +
         ' }');
@@ -170,9 +176,10 @@ Class.prototype.setMethod = function(name, args, body, flags) {
         );
     } else {
         // update the function
-        if (typeof flags !== 'undefined') method.setFlags(flags);
+        if (typeof flags !== 'undefined') method.setVisibility(flags);
         if (typeof args !== 'undefined') method.setArgs(args);
         if (typeof body !== 'undefined') method.setCode(body);
+        if (typeof returnType !== 'undefined') method.setReturnType(returnType);
     }
     return this;
 };
diff --git a/src/method.js b/src/method.js
index 1fcee5b..020078f 100644
--- a/src/method.js
+++ b/src/method.js
@@ -13,58 +13,98 @@ var editor = require('./helpers/editor');
  * @constructor
  */
 var Method = function Method(ast) {
-  this.ast = ast;
+    this.ast = ast;
 };
 editor(Method, 'body');
 
 /**
  * Change the method name
  */
-Method.prototype.setName = function(name) {
+Method.prototype.setName = function (name) {
     this.ast.name = name;
 };
 
 /**
  * Sets the abstract flag value
  */
-Method.prototype.setAbstract = function(flag) {
+Method.prototype.setAbstract = function (flag) {
     this.ast.isAbstract = flag === true;
     return this;
 };
 
 /**
- * Sets the abstract flag value
+ * Sets the final flag value
  */
-Method.prototype.setFinal = function(flag) {
+Method.prototype.setFinal = function (flag) {
     this.ast.isFinal = flag === true;
     return this;
 };
 
 /**
- * Sets the abstract flag value
+ * Sets the static flag value
  */
-Method.prototype.setStatic = function(flag) {
+Method.prototype.setStatic = function (flag) {
     this.ast.isStatic = flag === true;
     return this;
 };
 
 /**
- * Sets the abstract flag value
+ * Sets the visibility level value
  */
-Method.prototype.setVisibility = function(level) {
+Method.prototype.setVisibility = function (level) {
     this.ast.visibility = level;
     return this;
 };
 
 /**
- * Locate the node in the specified ast
+ * Sets the method arguments
+ */
+Method.prototype.setArgs = function (args) {
+    var ast = parser.parseEval('class a { \n' +
+        ' function dummy(' + args + ') { return null;}\n' +
+        ' }');
+
+    this.ast.arguments = ast.children[0].body[0].arguments;
+    return this;
+};
+
+/**
+ * Sets the method body
+ */
+Method.prototype.setCode = function (code) {
+    var ast = parser.parseEval('class a { \n' +
+        'public function dummy($dummy = true) { ' + code + '}\n' +
+        ' }');
+
+    this.ast.body = ast.children[0].body[0].body
+    return this;
+};
+
+/**
+ * Sets the return type
  */
-Method.locate = function(ast, name) {
-  return filter(ast, 'method', function(node) {
-    if (node.name === name) {
-      return new Method(node);
+Method.prototype.setReturnType = function (returnType) {
+    var ast = parser.parseEval('class a { \n' +
+        'public function dummy($dummy = true) ' + returnType + ' { return null; }\n' +
+        ' }');
+
+    this.ast.type = ast.children[0].body[0].type;
+
+    if (returnType[1] === '?') {
+        this.ast.nullable = true;
     }
-  });
+    return this;
+};
+
+/**
+ * Locate the node in the specified ast
+ */
+Method.locate = function (ast, name) {
+    return filter(ast, 'method', function (node) {
+        if (node.name === name) {
+            return new Method(node);
+        }
+    });
 };
 
 module.exports = Method;