Normalization of evolving unit vector in P2 function space

Dear Fenics community,
I would like to normalize an vector field \mathbf{n}(t) evolving with time t that is constrained to be
unit vector everywhere.
\frac{\mathrm{d}\mathbf{n}(t)}{\mathrm{d}t} = f(\mathbf{n}(t)), \quad |\mathbf{n}(t)|=1.
At this moment, the strategy I use to enforce the constraint is to normalize \mathbf{n} back to unit if it deviates too much.I know how to normalize n(t) if it lives in P1 function space. Something like

def normalize(w, dim=2, epsilon=0.01):                                                                                                                                                        
    """Check and normalize vector in P1 to norm 1."""                                                                                                                                   
    wv = w.vector()[:]                                                                                                                                                              
    wa = np.reshape(wv, (-1, dim))                                                                                                                                                  
    norms = np.linalg.norm(wa, axis=1, keepdims=True)                                                                                                                               
    if(np.max(np.abs(norms-1.)) > epsilon):                                                                                                                                         
        wa = wa/norms                                                                                                                                                               
        for i in range(dim):                                                                                                                                                        
            wv[i::dim] = wa[:, i]                                                                                                                                                   
        w.vector().set_local(wv)                                                                                                                                                    
        w.vector().apply('')                                                                                                                                                        
    return w  

will do the job.
But how to normalize n to unit vector if \mathbf{n} lives in P2 function space?
Or is there a better way to enforce the unit vector constraint?
Thanks in advance.

It seems to me that this should also work for P2 space.

Thank you for your reply.
I tried several examples. It indeed works. I am not familiar with how the degree of freedom is arranged in P2 and had the wrong impression that this only works for P1. Really should try before asking here. Thanks again.

No problem. The ordering of a vector FunctionSpace is always X^1, Y^1, Z^1, X^2, Y^2... no matter if the underlying scalar function space is P1 or P2 or something else