Writing code that is compatible across versions

Many Python modules these days are designed to run under multiple Python versions, especially supporting Python 2.x as well as Python 3.x versions. We will want to run the same tests in both versions, and to do this, we will need to write our code in such a way that the tests are compatible with both the versions.

Python's import mechanism gives us the flexibility we need to do this. At the top of the file, we import unittest like the following:

try:
   import unittest2 as unittest
except ImportError:
   import unittest

What this does is to first try and import unittest2. If we are running Python 2.x, then we should have installed this already. If it succeeds, then the module is imported and the module reference is renamed to unittest.

If we get an ImportError, then we are running Python 3.x, in which case we can import the unittest module bundled in the standard library.

Later in the code, we can just reference the unittest module and it will work normally.

This mechanism depends on the unittest2 module being always installed when using Python 2.x version. This is easily achieved by putting the unittest2 module as a dependency for only Python 2.x in our pip requirements file.

A similar approach works for mocks as follows:

try:
   from unittest import mock
except ImportError:
   import mock

Here we first try to import the mock library provided as a part of the unittest standard library module. This is available in Python 3.3 onward. If the import succeeds, then the mock library is imported. If it fails, it means that we are running an older Python version, so we directly import the mock library that we installed from PyPi.

Note how we use the line from unittest import mock instead of import unittest.mock. This is so that we end up with the same module reference name in both the cases. Once the import is done, we can reference the mock module in our code and it will work across Python versions.