The rule $\theta = \arctan \frac ba$ where $z = a+ib$ is wrong in multiple ways.
As you noticed, it does not correctly handle the case where $a = 0$ and $b \neq 0.$
It also gives incorrect answers for all numbers with negative real parts!
For example, let $z = -2 - i2.$ Then the arc tangent formula says
$$\theta = \arctan \frac {-2}{-2} = \arctan 1 = \frac\pi4,$$
but the correct answer is $\theta = \frac54\pi.$
In general, use the arc tangent formula only as a helpful hint to compute answers that are not obvious, and apply the obvious "fix" to the formula in the case where the real part is negative.
Remember, given $z = a + ib,$ the goal is to find $r$ and $\theta$ such that
\begin{align}
a &= r \cos\theta,\\
b &= r \sin\theta.
\end{align}
If a method achieves that result, use it; otherwise use a different method.