Skip to content

Commit

Permalink
Do not fix NamedTuple calls containing both a list of fields and ke…
Browse files Browse the repository at this point in the history
…ywords (astral-sh#5799)

## Summary

Fixes astral-sh#5794

## Test Plan

Existing tests
  • Loading branch information
harupy authored and evanrittenhouse committed Jul 19, 2023
1 parent d6dd698 commit 14d242f
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 130 deletions.
21 changes: 7 additions & 14 deletions crates/ruff/resources/test/fixtures/pyupgrade/UP014.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,9 @@
# with complex annotations
MyType = NamedTuple("MyType", [("a", int), ("b", tuple[str, ...])])

# with default values as list
MyType = NamedTuple(
"MyType",
[("a", int), ("b", str), ("c", list[bool])],
defaults=["foo", [True]],
)

# with namespace
MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])

# too many default values (OK)
MyType = NamedTuple(
"MyType",
[("a", int), ("b", str)],
defaults=[1, "bar", "baz"],
)

# invalid identifiers (OK)
MyType = NamedTuple("MyType", [("x-y", int), ("b", tuple[str, ...])])

Expand All @@ -29,3 +15,10 @@

# empty fields
MyType = typing.NamedTuple("MyType", [])

# keywords
MyType = typing.NamedTuple("MyType", a=int, b=tuple[str, ...])

# unfixable
MyType = typing.NamedTuple("MyType", [("a", int)], [("b", str)])
MyType = typing.NamedTuple("MyType", [("a", int)], b=str)
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,7 @@ fn match_named_tuple_assign<'a>(

/// Generate a `Stmt::AnnAssign` representing the provided property
/// definition.
fn create_property_assignment_stmt(
property: &str,
annotation: &Expr,
value: Option<&Expr>,
) -> Stmt {
fn create_property_assignment_stmt(property: &str, annotation: &Expr) -> Stmt {
ast::StmtAnnAssign {
target: Box::new(
ast::ExprName {
Expand All @@ -108,40 +104,15 @@ fn create_property_assignment_stmt(
.into(),
),
annotation: Box::new(annotation.clone()),
value: value.map(|value| Box::new(value.clone())),
value: None,
simple: true,
range: TextRange::default(),
}
.into()
}

/// Match the `defaults` keyword in a `NamedTuple(...)` call.
fn match_defaults(keywords: &[Keyword]) -> Result<&[Expr]> {
let defaults = keywords.iter().find(|keyword| {
if let Some(arg) = &keyword.arg {
arg == "defaults"
} else {
false
}
});
match defaults {
Some(defaults) => match &defaults.value {
Expr::List(ast::ExprList { elts, .. }) => Ok(elts),
Expr::Tuple(ast::ExprTuple { elts, .. }) => Ok(elts),
_ => bail!("Expected defaults to be `Expr::List` | `Expr::Tuple`"),
},
None => Ok(&[]),
}
}

/// Create a list of property assignments from the `NamedTuple` arguments.
fn create_properties_from_args(args: &[Expr], defaults: &[Expr]) -> Result<Vec<Stmt>> {
let Some(fields) = args.get(1) else {
let node = Stmt::Pass(ast::StmtPass {
range: TextRange::default(),
});
return Ok(vec![node]);
};
/// Create a list of property assignments from the `NamedTuple` fields argument.
fn create_properties_from_fields_arg(fields: &Expr) -> Result<Vec<Stmt>> {
let Expr::List(ast::ExprList { elts, .. }) = &fields else {
bail!("Expected argument to be `Expr::List`");
};
Expand All @@ -151,16 +122,8 @@ fn create_properties_from_args(args: &[Expr], defaults: &[Expr]) -> Result<Vec<S
});
return Ok(vec![node]);
}
let padded_defaults = if elts.len() >= defaults.len() {
std::iter::repeat(None)
.take(elts.len() - defaults.len())
.chain(defaults.iter().map(Some))
} else {
bail!("Defaults must be `None` or an iterable of at least the number of fields")
};
elts.iter()
.zip(padded_defaults)
.map(|(field, default)| {
.map(|field| {
let Expr::Tuple(ast::ExprTuple { elts, .. }) = &field else {
bail!("Expected `field` to be `Expr::Tuple`")
};
Expand All @@ -180,9 +143,21 @@ fn create_properties_from_args(args: &[Expr], defaults: &[Expr]) -> Result<Vec<S
if is_dunder(property) {
bail!("Cannot use dunder property name: {}", property)
}
Ok(create_property_assignment_stmt(
property, annotation, default,
))
Ok(create_property_assignment_stmt(property, annotation))
})
.collect()
}

/// Create a list of property assignments from the `NamedTuple` keyword arguments.
fn create_properties_from_keywords(keywords: &[Keyword]) -> Result<Vec<Stmt>> {
keywords
.iter()
.map(|keyword| {
let Keyword { arg, value, .. } = keyword;
let Some(arg) = arg else {
bail!("Expected `keyword` to have an `arg`")
};
Ok(create_property_assignment_stmt(arg.as_str(), value))
})
.collect()
}
Expand Down Expand Up @@ -228,12 +203,32 @@ pub(crate) fn convert_named_tuple_functional_to_class(
return;
};

let properties = match match_defaults(keywords)
.and_then(|defaults| create_properties_from_args(args, defaults))
{
Ok(properties) => properties,
Err(err) => {
debug!("Skipping `NamedTuple` \"{typename}\": {err}");
let properties = match (&args[1..], keywords) {
// Ex) NamedTuple("MyType")
([], []) => vec![Stmt::Pass(ast::StmtPass {
range: TextRange::default(),
})],
// Ex) NamedTuple("MyType", [("a", int), ("b", str)])
([fields], []) => {
if let Ok(properties) = create_properties_from_fields_arg(fields) {
properties
} else {
debug!("Skipping `NamedTuple` \"{typename}\": unable to parse fields");
return;
}
}
// Ex) NamedTuple("MyType", a=int, b=str)
([], keywords) => {
if let Ok(properties) = create_properties_from_keywords(keywords) {
properties
} else {
debug!("Skipping `NamedTuple` \"{typename}\": unable to parse keywords");
return;
}
}
// Unfixable
_ => {
debug!("Skipping `NamedTuple` \"{typename}\": mixed fields and keywords");
return;
}
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ UP014.py:5:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class s
5 | MyType = NamedTuple("MyType", [("a", int), ("b", tuple[str, ...])])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
6 |
7 | # with default values as list
7 | # with namespace
|
= help: Convert `MyType` to class syntax

Expand All @@ -20,97 +20,93 @@ UP014.py:5:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class s
6 |+ a: int
7 |+ b: tuple[str, ...]
6 8 |
7 9 | # with default values as list
8 10 | MyType = NamedTuple(
7 9 | # with namespace
8 10 | MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])

UP014.py:8:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
|
7 | # with default values as list
8 | / MyType = NamedTuple(
9 | | "MyType",
10 | | [("a", int), ("b", str), ("c", list[bool])],
11 | | defaults=["foo", [True]],
12 | | )
| |_^ UP014
13 |
14 | # with namespace
7 | # with namespace
8 | MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
9 |
10 | # invalid identifiers (OK)
|
= help: Convert `MyType` to class syntax

Suggested fix
5 5 | MyType = NamedTuple("MyType", [("a", int), ("b", tuple[str, ...])])
6 6 |
7 7 | # with default values as list
8 |-MyType = NamedTuple(
9 |- "MyType",
10 |- [("a", int), ("b", str), ("c", list[bool])],
11 |- defaults=["foo", [True]],
12 |-)
8 |+class MyType(NamedTuple):
7 7 | # with namespace
8 |-MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])
8 |+class MyType(typing.NamedTuple):
9 |+ a: int
10 |+ b: str = "foo"
11 |+ c: list[bool] = [True]
13 12 |
14 13 | # with namespace
15 14 | MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])
10 |+ b: str
9 11 |
10 12 | # invalid identifiers (OK)
11 13 | MyType = NamedTuple("MyType", [("x-y", int), ("b", tuple[str, ...])])

UP014.py:15:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
UP014.py:14:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
|
14 | # with namespace
15 | MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
16 |
17 | # too many default values (OK)
13 | # no fields
14 | MyType = typing.NamedTuple("MyType")
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
15 |
16 | # empty fields
|
= help: Convert `MyType` to class syntax

Suggested fix
12 12 | )
13 13 |
14 14 | # with namespace
15 |-MyType = typing.NamedTuple("MyType", [("a", int), ("b", str)])
15 |+class MyType(typing.NamedTuple):
16 |+ a: int
17 |+ b: str
16 18 |
17 19 | # too many default values (OK)
18 20 | MyType = NamedTuple(
11 11 | MyType = NamedTuple("MyType", [("x-y", int), ("b", tuple[str, ...])])
12 12 |
13 13 | # no fields
14 |-MyType = typing.NamedTuple("MyType")
14 |+class MyType(typing.NamedTuple):
15 |+ pass
15 16 |
16 17 | # empty fields
17 18 | MyType = typing.NamedTuple("MyType", [])

UP014.py:28:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
UP014.py:17:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
|
27 | # no fields
28 | MyType = typing.NamedTuple("MyType")
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
29 |
30 | # empty fields
16 | # empty fields
17 | MyType = typing.NamedTuple("MyType", [])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
18 |
19 | # keywords
|
= help: Convert `MyType` to class syntax

Suggested fix
25 25 | MyType = NamedTuple("MyType", [("x-y", int), ("b", tuple[str, ...])])
26 26 |
27 27 | # no fields
28 |-MyType = typing.NamedTuple("MyType")
28 |+class MyType(typing.NamedTuple):
29 |+ pass
29 30 |
30 31 | # empty fields
31 32 | MyType = typing.NamedTuple("MyType", [])
14 14 | MyType = typing.NamedTuple("MyType")
15 15 |
16 16 | # empty fields
17 |-MyType = typing.NamedTuple("MyType", [])
17 |+class MyType(typing.NamedTuple):
18 |+ pass
18 19 |
19 20 | # keywords
20 21 | MyType = typing.NamedTuple("MyType", a=int, b=tuple[str, ...])

UP014.py:31:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
UP014.py:20:1: UP014 [*] Convert `MyType` from `NamedTuple` functional to class syntax
|
30 | # empty fields
31 | MyType = typing.NamedTuple("MyType", [])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
19 | # keywords
20 | MyType = typing.NamedTuple("MyType", a=int, b=tuple[str, ...])
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ UP014
21 |
22 | # unfixable
|
= help: Convert `MyType` to class syntax

Suggested fix
28 28 | MyType = typing.NamedTuple("MyType")
29 29 |
30 30 | # empty fields
31 |-MyType = typing.NamedTuple("MyType", [])
31 |+class MyType(typing.NamedTuple):
32 |+ pass
17 17 | MyType = typing.NamedTuple("MyType", [])
18 18 |
19 19 | # keywords
20 |-MyType = typing.NamedTuple("MyType", a=int, b=tuple[str, ...])
20 |+class MyType(typing.NamedTuple):
21 |+ a: int
22 |+ b: tuple[str, ...]
21 23 |
22 24 | # unfixable
23 25 | MyType = typing.NamedTuple("MyType", [("a", int)], [("b", str)])


0 comments on commit 14d242f

Please sign in to comment.