Format import statements (#5493)

## Summary Format import statements in all their variants. Specifically, this implemented formatting `StmtImport`, `StmtImportFrom` and `Alias`. ## Test Plan I added some custom snapshots, even though this has been covered well by black's tests.
astral-sh · Jul 4, 2023 · 787e2fd · 787e2fd
1 parent 6acc316
commit 787e2fd
Show file tree

Hide file tree

Showing 26 changed files with 308 additions and 995 deletions.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import.py
@@ -0,0 +1,3 @@
+from a import aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa
+from a import aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa, aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa
+from a import aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as dfgsdfgsd, aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/import_from.py
@@ -0,0 +1,16 @@
+from a import aksjdhflsakhdflkjsadlfajkslhf
+from a import (
+    aksjdhflsakhdflkjsadlfajkslhf,
+)
+from a import (
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa,
+)
+from a import (
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa,
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa,
+)
+from a import (
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as dfgsdfgsd,
+    aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa as sdkjflsdjlahlfd,
+)
+from aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa import *
diff --git a/crates/ruff_python_formatter/src/other/alias.rs b/crates/ruff_python_formatter/src/other/alias.rs
@@ -1,12 +1,22 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{space, text};
+use ruff_formatter::{write, Buffer, Format, FormatResult};
 use rustpython_parser::ast::Alias;
 
 #[derive(Default)]
 pub struct FormatAlias;
 
 impl FormatNodeRule<Alias> for FormatAlias {
     fn fmt_fields(&self, item: &Alias, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let Alias {
+            range: _,
+            name,
+            asname,
+        } = item;
+        name.format().fmt(f)?;
+        if let Some(asname) = asname {
+            write!(f, [space(), text("as"), space(), asname.format()])?;
+        }
+        Ok(())
     }
 }
diff --git a/crates/ruff_python_formatter/src/statement/stmt_import.rs b/crates/ruff_python_formatter/src/statement/stmt_import.rs
@@ -1,4 +1,5 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
+use crate::{FormatNodeRule, FormattedIterExt, PyFormatter};
+use ruff_formatter::prelude::{format_args, format_with, space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
 use rustpython_parser::ast::StmtImport;
 
@@ -7,6 +8,12 @@ pub struct FormatStmtImport;
 
 impl FormatNodeRule<StmtImport> for FormatStmtImport {
     fn fmt_fields(&self, item: &StmtImport, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let StmtImport { names, range: _ } = item;
+        let names = format_with(|f| {
+            f.join_with(&format_args![text(","), space()])
+                .entries(names.iter().formatted())
+                .finish()
+        });
+        write!(f, [text("import"), space(), names])
     }
 }
diff --git a/crates/ruff_python_formatter/src/statement/stmt_import_from.rs b/crates/ruff_python_formatter/src/statement/stmt_import_from.rs
@@ -1,12 +1,48 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::builders::{optional_parentheses, PyFormatterExtensions};
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{dynamic_text, format_with, space, text};
+use ruff_formatter::{write, Buffer, Format, FormatResult};
 use rustpython_parser::ast::StmtImportFrom;
 
 #[derive(Default)]
 pub struct FormatStmtImportFrom;
 
 impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
     fn fmt_fields(&self, item: &StmtImportFrom, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let StmtImportFrom {
+            module,
+            names,
+            range: _,
+            level,
+        } = item;
+
+        let level_str = level
+            .map(|level| ".".repeat(level.to_usize()))
+            .unwrap_or(String::default());
+
+        write!(
+            f,
+            [
+                text("from"),
+                space(),
+                dynamic_text(&level_str, None),
+                module.as_ref().map(AsFormat::format),
+                space(),
+                text("import"),
+                space(),
+            ]
+        )?;
+        if let [name] = names.as_slice() {
+            // star can't be surrounded by parentheses
+            if name.name.as_str() == "*" {
+                return text("*").fmt(f);
+            }
+        }
+        let names = format_with(|f| {
+            f.join_comma_separated()
+                .entries(names.iter().map(|name| (name, name.format())))
+                .finish()
+        });
+        optional_parentheses(&names).fmt(f)
     }
 }
diff --git a/...uff_python_formatter/tests/snapshots/black_compatibility@miscellaneous__force_pyi.py.snap b/...uff_python_formatter/tests/snapshots/black_compatibility@miscellaneous__force_pyi.py.snap
@@ -43,21 +43,20 @@ def eggs() -> Union[str, int]: ...
 --- Black
 +++ Ruff
 @@ -1,32 +1,58 @@
--from typing import Union
-+NOT_YET_IMPLEMENTED_StmtImportFrom
-+
+ from typing import Union
  
++
  @bird
 -def zoo(): ...
 +def zoo():
 +    ...
-+
-+
-+class A:
-+    ...
  
 -class A: ...
  
++class A:
++    ...
++
++
  @bar
  class B:
 -    def BMethod(self) -> None: ...
@@ -94,14 +93,14 @@ def eggs() -> Union[str, int]: ...
 +
 +class F(A, C):
 +    ...
++
++
++def spam() -> None:
++    ...
  
 -class F(A, C): ...
  
 -def spam() -> None: ...
-+def spam() -> None:
-+    ...
-+
-+
  @overload
 -def spam(arg: str) -> str: ...
 +def spam(arg: str) -> str:
@@ -120,7 +119,7 @@ def eggs() -> Union[str, int]: ...
 ## Ruff Output
 
 ```py
-NOT_YET_IMPLEMENTED_StmtImportFrom
+from typing import Union
 
 
 @bird

diff --git a/...hon_formatter/tests/snapshots/black_compatibility@py_310__pattern_matching_extras.py.snap b/...hon_formatter/tests/snapshots/black_compatibility@py_310__pattern_matching_extras.py.snap
@@ -132,8 +132,7 @@ match bar1:
 --- Black
 +++ Ruff
 @@ -1,119 +1,43 @@
--import match
-+NOT_YET_IMPLEMENTED_StmtImport
+ import match
  
 -match something:
 -    case [a as b]:
@@ -208,10 +207,11 @@ match bar1:
 -        ),
 -    ):
 -        pass
--
++NOT_YET_IMPLEMENTED_StmtMatch
+ 
 -    case [a as match]:
 -        pass
--
+ 
 -    case case:
 -        pass
 +NOT_YET_IMPLEMENTED_StmtMatch
@@ -220,8 +220,9 @@ match bar1:
 -match match:
 -    case case:
 -        pass
--
--
++NOT_YET_IMPLEMENTED_StmtMatch
+ 
+ 
 -match a, *b(), c:
 -    case d, *f, g:
 -        pass
@@ -236,30 +237,28 @@ match bar1:
 -        pass
 -    case {"maybe": something(complicated as this) as that}:
 -        pass
+-
 +NOT_YET_IMPLEMENTED_StmtMatch
  
- 
 -match something:
 -    case 1 as a:
 -        pass
-+NOT_YET_IMPLEMENTED_StmtMatch
  
 -    case 2 as b, 3 as c:
 -        pass
++NOT_YET_IMPLEMENTED_StmtMatch
  
 -    case 4 as d, (5 as e), (6 | 7 as g), *h:
 -        pass
-+NOT_YET_IMPLEMENTED_StmtMatch
- 
  
+-
 -match bar1:
 -    case Foo(aa=Callable() as aa, bb=int()):
 -        print(bar1.aa, bar1.bb)
 -    case _:
 -        print("no match", "\n")
-+NOT_YET_IMPLEMENTED_StmtMatch
- 
- 
+-
+-
 -match bar1:
 -    case Foo(
 -        normal=x, perhaps=[list, {"x": d, "y": 1.0}] as y, otherwise=something, q=t as u
@@ -271,7 +270,7 @@ match bar1:
 ## Ruff Output
 
 ```py
-NOT_YET_IMPLEMENTED_StmtImport
+import match
 
 NOT_YET_IMPLEMENTED_StmtMatch
 

diff --git a/...ff_python_formatter/tests/snapshots/black_compatibility@simple_cases__collections.py.snap b/...ff_python_formatter/tests/snapshots/black_compatibility@simple_cases__collections.py.snap
@@ -83,31 +83,9 @@ if True:
 ```diff
 --- Black
 +++ Ruff
-@@ -1,40 +1,22 @@
--import core, time, a
-+NOT_YET_IMPLEMENTED_StmtImport
- 
--from . import A, B, C
-+NOT_YET_IMPLEMENTED_StmtImportFrom
- 
- # keeps existing trailing comma
--from foo import (
--    bar,
--)
-+NOT_YET_IMPLEMENTED_StmtImportFrom
- 
- # also keeps existing structure
--from foo import (
--    baz,
--    qux,
--)
-+NOT_YET_IMPLEMENTED_StmtImportFrom
- 
- # `as` works as well
--from foo import (
--    xyzzy as magic,
--)
-+NOT_YET_IMPLEMENTED_StmtImportFrom
+@@ -18,23 +18,12 @@
+     xyzzy as magic,
+ )
  
 -a = {
 -    1,
@@ -132,7 +110,7 @@ if True:
  nested_no_trailing_comma = {(1, 2, 3), (4, 5, 6)}
  nested_long_lines = [
      "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
-@@ -52,10 +34,7 @@
+@@ -52,10 +41,7 @@
  y = {
      "oneple": (1,),
  }
@@ -149,18 +127,25 @@ if True:
 ## Ruff Output
 
 ```py
-NOT_YET_IMPLEMENTED_StmtImport
+import core, time, a
 
-NOT_YET_IMPLEMENTED_StmtImportFrom
+from . import A, B, C
 
 # keeps existing trailing comma
-NOT_YET_IMPLEMENTED_StmtImportFrom
+from foo import (
+    bar,
+)
 
 # also keeps existing structure
-NOT_YET_IMPLEMENTED_StmtImportFrom
+from foo import (
+    baz,
+    qux,
+)
 
 # `as` works as well
-NOT_YET_IMPLEMENTED_StmtImportFrom
+from foo import (
+    xyzzy as magic,
+)
 
 a = {1, 2, 3}
 b = {1, 2, 3}