Missing import.
Replacing use of unitialised attribute with method param.
Merge pull request #166 from william-ferguson-au/missing-import

Fixes for pull requests applied earlier.