Re-organize code to resolve docstring bug

[ablaom]

Resolves #33

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.58%. Comparing base (21fce28) to head (89c8479).
⚠️ Report is 2 commits behind head on dev.

Additional details and impacted files

@@            Coverage Diff             @@
##              dev      #37      +/-   ##
==========================================
+ Coverage   92.94%   93.58%   +0.64%     
==========================================
  Files           4        4              
  Lines         156      156              
==========================================
+ Hits          145      146       +1     
+ Misses         11       10       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.