Minor Improvement in atoms() and Tests

Author: ShikharJ

symengine/lib/symengine_wrapper.pyx

@@ -585,12 +585,15 @@ cdef class Basic(object):
         return ring(self._sage_())
 
     def atoms(self, *types):
-        s = set()
-        if (isinstance(self, types)):
-            s.add(self)
-        for arg in self.args:
-            s.update(arg.atoms(*types))
-        return s
+        if types:
+            s = set()
+            if (isinstance(self, types)):
+                s.add(self)
+            for arg in self.args:
+                s.update(arg.atoms(*types))
+            return s
+        else:
+            return self.free_symbols
 
     def simplify(self, *args, **kwargs):
         return sympify(self._sympy_().simplify(*args, **kwargs))
@@ -1969,12 +1972,15 @@ cdef class DenseMatrix(MatrixBase):
         return self[:]
 
     def atoms(self, *types):
-        s = set()
-        if (isinstance(self, types)):
-            s.add(self)
-        for arg in self.tolist():
-            s.update(arg.atoms(*types))
-        return s
+        if types:
+            s = set()
+            if (isinstance(self, types)):
+                s.add(self)
+            for arg in self.tolist():
+                s.update(arg.atoms(*types))
+            return s
+        else:
+           return self.free_symbols
 
     def simplify(self, *args, **kwargs):
         return self._applyfunc(lambda x : x.simplify(*args, **kwargs))

symengine/tests/test_arit.py

@@ -119,6 +119,15 @@ def test_args():
     assert (2*x**2).args == (2, x**2)
     assert set((2*x**2*y).args) == set((Integer(2), x**2, y))
 
+def test_atoms():
+    x = Symbol("x")
+    y = Symbol("y")
+    z = Symbol("z")
+    assert (x**2).atoms() == set([x])
+    assert (x**2).atoms(Symbol) == set([x])
+    assert (x ** y + z).atoms() == set([x, y, z])
+    assert (x**y + z).atoms(Symbol) == set([x, y, z])
+
 def test_free_symbols():
     x = Symbol("x")
     y = Symbol("y")