def modify_node(self, node):
"""Modifies the ``ast.Call`` node to call ``TrainingInput`` instead.
Args:
node (ast.Call): a node that represents a *TrainingInput constructor.
"""
if matching.matches_name(node, S3_INPUT_NAME):
node.func.id = "TrainingInput"
elif matching.matches_attr(node, S3_INPUT_NAME):
node.func.attr = "TrainingInput"
_rename_namespace(node, "session")
return node
def _rename_class(node):
"""Renames the RealTimePredictor base class to Predictor"""
if matching.matches_name(node, BASE_PREDICTOR):
node.func.id = "Predictor"
elif matching.matches_attr(node, BASE_PREDICTOR):
node.func.attr = "Predictor"
def test_matches_name():
assert matching.matches_name(ast_call("KMeans()"), "KMeans")
assert not matching.matches_name(ast_call("sagemaker.KMeans()"), "KMeans")
assert not matching.matches_name(ast_call("MXNet()"), "KMeans")