Production of heavy flavoured hadrons can serve as a probe of QCD processes and, in the case of production in heavy ion collisions, also as a probe of the properties of deconfined quark-gluon plasma. This paper describes the latest results in the production of charmonium in the ATLAS experiment at the Large Hadron Collider. The presented analyses include productions of $J/\psi$ and $\psi$(2S) in $pp$ collisions and $\sqrt{s}$ = 7, 8 and 13 TeV, and production of $J/\psi$, $\psi$(2S) and $\Upsilon$(nS) in $pp$ and $p$+Pb collisions at $\sqrt{s}$ = 5.02 TeV. In both the analyses prompt and non-prompt productions are disentangled. In the latter analysis production suppression in $p$+Pb with respect to $pp$ collisions is measured and emphasis is put on the collision centrality dependence.