Th problem with square root for matrices is the lack of uniqueness, you don't have an integral ring, hence you can have infinite roots for the polynomial $X^2 -A$. For example $X^2 = 0$ in $\mathcal{M}_2(\mathbb{R})$ has all nilpotent matrices as solution (and $0$ is diagonalisable).
One way to achieve uniqueness is to work with symmetric definite positive matrices, which have an unique square root.
Otherwise you could use for $||| A-I||| < 1$:
$$A^{1/2} = (I + (A-I))^{1/2} = \sum_{k\geq 0} \binom{1/2}{k}A^k.$$
This gives you a (not unique) but canonical square root.
You solution would work and would be a canonical way to find a square root, to detail it, it is better to work with the eigenspaces, if you have
$\lambda_1,\dots,\lambda_d$ as eigenvalues, you define $A_{|E_\lambda} = \lambda^{1/2} id$, but think about it, the uniqueness of the square root is not even verified in $\mathbb{C}$ and if we were to choose a canonical root, we would have to do it on $\mathbb{C} \setminus \mathbb{R}_-$, so you have to ensure something on the eigenvalues.
One last solution is to work with the D+N decomposition, (or dunford decomposition) in case you don't have a diagonalisable matrix but you won't find a canonical way to find it.
If you think about it, you would very much like to find a square root that is a polynomial in $A$ so finding a root in $k[A]$ would be equivalent to finding a root of $\bar{X}$ in $k[X]/(\pi_A)$