CallPath newtype wrapper (#10201)

## Summary

This PR changes the `CallPath` type alias to a newtype wrapper. 

A newtype wrapper allows us to limit the API and to experiment with
alternative ways to implement matching on `CallPath`s.



## Test Plan

`cargo test`
This commit is contained in:
Micha Reiser 2024-03-03 16:54:24 +01:00 committed by GitHub
parent fb05d218c3
commit e725b6fdaf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
165 changed files with 551 additions and 433 deletions

View file

@ -1,19 +1,154 @@
use smallvec::{smallvec, SmallVec};
use smallvec::SmallVec;
use std::fmt::{Display, Formatter, Write};
use crate::{nodes, Expr};
/// A representation of a qualified name, like `typing.List`.
pub type CallPath<'a> = SmallVec<[&'a str; 8]>;
#[derive(Debug, Clone, Eq, Hash)]
pub struct CallPath<'a> {
segments: SmallVec<[&'a str; 8]>,
}
impl<'a> CallPath<'a> {
pub fn from_expr(expr: &'a Expr) -> Option<Self> {
let segments = collect_call_path(expr)?;
Some(Self { segments })
}
/// Create a [`CallPath`] from an unqualified name.
///
/// ```rust
/// # use smallvec::smallvec;
/// # use ruff_python_ast::call_path::CallPath;
///
/// assert_eq!(CallPath::from_unqualified_name("typing.List").segments(), ["typing", "List"]);
/// assert_eq!(CallPath::from_unqualified_name("list").segments(), ["list"]);
/// ```
#[inline]
pub fn from_unqualified_name(name: &'a str) -> Self {
name.split('.').collect()
}
/// Create a [`CallPath`] from a fully-qualified name.
///
/// ```rust
/// # use smallvec::smallvec;
/// # use ruff_python_ast::call_path::CallPath;
///
/// assert_eq!(CallPath::from_qualified_name("typing.List").segments(), ["typing", "List"]);
/// assert_eq!(CallPath::from_qualified_name("list").segments(), ["", "list"]);
/// ```
#[inline]
pub fn from_qualified_name(name: &'a str) -> Self {
if let Some(dot) = name.find('.') {
let mut segments = SmallVec::new();
segments.push(&name[..dot]);
segments.extend(name[dot + 1..].split('.'));
Self { segments }
} else {
// Special-case: for builtins, return `["", "int"]` instead of `["int"]`.
Self::from_slice(&["", name])
}
}
#[inline]
pub fn from_slice(segments: &[&'a str]) -> Self {
Self {
segments: segments.into(),
}
}
pub fn starts_with(&self, other: &CallPath) -> bool {
self.segments().starts_with(other.segments())
}
#[inline]
pub fn segments(&self) -> &[&'a str] {
&self.segments
}
#[inline]
pub fn into_boxed_slice(self) -> Box<[&'a str]> {
self.segments.into_boxed_slice()
}
}
impl<'a> FromIterator<&'a str> for CallPath<'a> {
fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
Self {
segments: iter.into_iter().collect(),
}
}
}
impl<'a, 'b> PartialEq<CallPath<'b>> for CallPath<'a> {
#[inline]
fn eq(&self, other: &CallPath<'b>) -> bool {
self.segments == other.segments
}
}
#[derive(Debug, Clone, Default)]
pub struct CallPathBuilder<'a> {
segments: SmallVec<[&'a str; 8]>,
}
impl<'a> CallPathBuilder<'a> {
pub fn with_capacity(capacity: usize) -> Self {
Self {
segments: SmallVec::with_capacity(capacity),
}
}
pub fn new() -> Self {
Self::default()
}
pub fn from_path(call_path: CallPath<'a>) -> Self {
Self {
segments: call_path.segments,
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.segments.is_empty()
}
#[inline]
pub fn push(&mut self, segment: &'a str) {
self.segments.push(segment);
}
pub fn pop(&mut self) {
self.segments.pop();
}
#[inline]
pub fn extend(&mut self, segments: impl IntoIterator<Item = &'a str>) {
self.segments.extend(segments);
}
pub fn extend_from_slice(&mut self, segments: &[&'a str]) {
self.segments.extend_from_slice(segments);
}
pub fn build(self) -> CallPath<'a> {
CallPath {
segments: self.segments,
}
}
}
/// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`).
pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
fn collect_call_path(expr: &Expr) -> Option<SmallVec<[&str; 8]>> {
// Unroll the loop up to eight times, to match the maximum number of expected attributes.
// In practice, unrolling appears to give about a 4x speed-up on this hot path.
let attr1 = match expr {
Expr::Attribute(attr1) => attr1,
// Ex) `foo`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[id.as_str()]))
return Some(SmallVec::from_slice(&[id.as_str()]))
}
_ => return None,
};
@ -22,7 +157,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr2) => attr2,
// Ex) `foo.bar`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[id.as_str(), attr1.attr.as_str()]))
return Some(SmallVec::from_slice(&[id.as_str(), attr1.attr.as_str()]))
}
_ => return None,
};
@ -31,7 +166,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr3) => attr3,
// Ex) `foo.bar.baz`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from_slice(&[
id.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
@ -44,7 +179,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr4) => attr4,
// Ex) `foo.bar.baz.bop`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from_slice(&[
id.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
@ -58,7 +193,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr5) => attr5,
// Ex) `foo.bar.baz.bop.bap`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from_slice(&[
id.as_str(),
attr4.attr.as_str(),
attr3.attr.as_str(),
@ -73,7 +208,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr6) => attr6,
// Ex) `foo.bar.baz.bop.bap.bab`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from_slice(&[
id.as_str(),
attr5.attr.as_str(),
attr4.attr.as_str(),
@ -89,7 +224,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr7) => attr7,
// Ex) `foo.bar.baz.bop.bap.bab.bob`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from_slice(&[
id.as_str(),
attr6.attr.as_str(),
attr5.attr.as_str(),
@ -106,7 +241,7 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
Expr::Attribute(attr8) => attr8,
// Ex) `foo.bar.baz.bop.bap.bab.bob.bib`
Expr::Name(nodes::ExprName { id, .. }) => {
return Some(CallPath::from_slice(&[
return Some(SmallVec::from([
id.as_str(),
attr7.attr.as_str(),
attr6.attr.as_str(),
@ -135,70 +270,59 @@ pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
})
}
/// Convert an `Expr` to its call path (like `List`, or `typing.List`).
pub fn compose_call_path(expr: &Expr) -> Option<String> {
collect_call_path(expr).map(|call_path| format_call_path(&call_path))
impl Display for CallPath<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
format_call_path_segments(self.segments(), f)
}
}
/// Format a call path for display.
pub fn format_call_path(call_path: &[&str]) -> String {
if call_path.first().map_or(false, |first| first.is_empty()) {
/// Convert an `Expr` to its call path (like `List`, or `typing.List`).
pub fn compose_call_path(expr: &Expr) -> Option<String> {
CallPath::from_expr(expr).map(|call_path| call_path.to_string())
}
pub fn format_call_path_segments(segments: &[&str], w: &mut dyn Write) -> std::fmt::Result {
if segments.first().is_some_and(|first| first.is_empty()) {
// If the first segment is empty, the `CallPath` is that of a builtin.
// Ex) `["", "bool"]` -> `"bool"`
call_path[1..].join(".")
} else if call_path
.first()
.map_or(false, |first| matches!(*first, "."))
{
let mut first = true;
for segment in segments.iter().skip(1) {
if !first {
w.write_char('.')?;
}
w.write_str(segment)?;
first = false;
}
} else if segments.first().is_some_and(|first| matches!(*first, ".")) {
// If the call path is dot-prefixed, it's an unresolved relative import.
// Ex) `[".foo", "bar"]` -> `".foo.bar"`
let mut formatted = String::new();
let mut iter = call_path.iter();
let mut iter = segments.iter();
for segment in iter.by_ref() {
if *segment == "." {
formatted.push('.');
w.write_char('.')?;
} else {
formatted.push_str(segment);
w.write_str(segment)?;
break;
}
}
for segment in iter {
formatted.push('.');
formatted.push_str(segment);
w.write_char('.')?;
w.write_str(segment)?;
}
formatted
} else {
call_path.join(".")
}
}
let mut first = true;
for segment in segments {
if !first {
w.write_char('.')?;
}
/// Create a [`CallPath`] from an unqualified name.
///
/// ```rust
/// # use smallvec::smallvec;
/// # use ruff_python_ast::call_path::from_unqualified_name;
///
/// assert_eq!(from_unqualified_name("typing.List").as_slice(), ["typing", "List"]);
/// assert_eq!(from_unqualified_name("list").as_slice(), ["list"]);
/// ```
pub fn from_unqualified_name(name: &str) -> CallPath {
name.split('.').collect()
}
/// Create a [`CallPath`] from a fully-qualified name.
///
/// ```rust
/// # use smallvec::smallvec;
/// # use ruff_python_ast::call_path::from_qualified_name;
///
/// assert_eq!(from_qualified_name("typing.List").as_slice(), ["typing", "List"]);
/// assert_eq!(from_qualified_name("list").as_slice(), ["", "list"]);
/// ```
pub fn from_qualified_name(name: &str) -> CallPath {
if name.contains('.') {
name.split('.').collect()
} else {
// Special-case: for builtins, return `["", "int"]` instead of `["int"]`.
smallvec!["", name]
w.write_str(segment)?;
first = false;
}
}
Ok(())
}

View file

@ -2,13 +2,12 @@ use std::borrow::Cow;
use std::path::Path;
use rustc_hash::FxHashMap;
use smallvec::SmallVec;
use ruff_python_trivia::{indentation_at_offset, CommentRanges, SimpleTokenKind, SimpleTokenizer};
use ruff_source_file::Locator;
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
use crate::call_path::CallPath;
use crate::call_path::{CallPath, CallPathBuilder};
use crate::parenthesize::parenthesized_range;
use crate::statement_visitor::StatementVisitor;
use crate::visitor::Visitor;
@ -793,16 +792,16 @@ pub fn to_module_path(package: &Path, path: &Path) -> Option<Vec<String>> {
/// ```rust
/// # use ruff_python_ast::helpers::collect_import_from_member;
///
/// assert_eq!(collect_import_from_member(None, None, "bar").as_slice(), ["bar"]);
/// assert_eq!(collect_import_from_member(Some(1), None, "bar").as_slice(), [".", "bar"]);
/// assert_eq!(collect_import_from_member(Some(1), Some("foo"), "bar").as_slice(), [".", "foo", "bar"]);
/// assert_eq!(collect_import_from_member(None, None, "bar").segments(), ["bar"]);
/// assert_eq!(collect_import_from_member(Some(1), None, "bar").segments(), [".", "bar"]);
/// assert_eq!(collect_import_from_member(Some(1), Some("foo"), "bar").segments(), [".", "foo", "bar"]);
/// ```
pub fn collect_import_from_member<'a>(
level: Option<u32>,
module: Option<&'a str>,
member: &'a str,
) -> CallPath<'a> {
let mut call_path: CallPath = SmallVec::with_capacity(
let mut call_path_builder = CallPathBuilder::with_capacity(
level.unwrap_or_default() as usize
+ module
.map(|module| module.split('.').count())
@ -814,20 +813,20 @@ pub fn collect_import_from_member<'a>(
if let Some(level) = level {
if level > 0 {
for _ in 0..level {
call_path.push(".");
call_path_builder.push(".");
}
}
}
// Add the remaining segments.
if let Some(module) = module {
call_path.extend(module.split('.'));
call_path_builder.extend(module.split('.'));
}
// Add the member.
call_path.push(member);
call_path_builder.push(member);
call_path
call_path_builder.build()
}
/// Format the call path for a relative import, or `None` if the relative import extends beyond
@ -840,27 +839,28 @@ pub fn from_relative_import<'a>(
// The remaining segments to the call path (e.g., given `bar.baz`, `["baz"]`).
tail: &[&'a str],
) -> Option<CallPath<'a>> {
let mut call_path: CallPath = SmallVec::with_capacity(module.len() + import.len() + tail.len());
let mut call_path_builder =
CallPathBuilder::with_capacity(module.len() + import.len() + tail.len());
// Start with the module path.
call_path.extend(module.iter().map(String::as_str));
call_path_builder.extend(module.iter().map(String::as_str));
// Remove segments based on the number of dots.
for segment in import {
if *segment == "." {
if call_path.is_empty() {
if call_path_builder.is_empty() {
return None;
}
call_path.pop();
call_path_builder.pop();
} else {
call_path.push(segment);
call_path_builder.push(segment);
}
}
// Add the remaining segments.
call_path.extend_from_slice(tail);
call_path_builder.extend_from_slice(tail);
Some(call_path)
Some(call_path_builder.build())
}
/// Given an imported module (based on its relative import level and module name), return the