Refactor method handling to be more flexible