Merge pull request #81 from project-eutopia/replace_method_can_take_s…

…trings

replace method accepts ::String and AST::Variable

lib/keisan/ast/variable.rb

@@ -60,11 +60,8 @@ def simplify(context = nil)
       end
 
       def replace(variable, replacement)
-        if name == variable.name
-          replacement
-        else
-          self
-        end
+        to_replace_name = variable.is_a?(::String) ? variable : variable.name
+        name == to_replace_name ? replacement : self
       end
 
       def differentiate(variable, context = nil)

spec/keisan/ast/node_spec.rb

@@ -658,6 +658,14 @@
       ast = Keisan::AST.parse("replace(x**2 + 1 / x, x, 10)")
       expect(ast.value).to eq (100 + Rational(1,10))
     end
+
+    it "accepts string or variable as argument" do
+      ast1 = Keisan::AST.parse("x**2").replace(Keisan::AST::Variable.new("x"), 3)
+      ast2 = Keisan::AST.parse("x**2").replace("x", 5)
+
+      expect(ast1.to_s).to eq("3**2")
+      expect(ast2.to_s).to eq("5**2")
+    end
   end
 
   describe "multi line" do