Skip to content

Commit a697a25

Browse files
committed
REVIEW
1 parent 45c9059 commit a697a25

File tree

4 files changed

+96
-48
lines changed

4 files changed

+96
-48
lines changed

pyrefly/lib/alt/function.rs

Lines changed: 72 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -65,27 +65,86 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
6565
|| class_metadata.is_some_and(|idx| self.get_idx(*idx).is_protocol());
6666
let def = self.get_idx(idx);
6767
if def.metadata.flags.is_overload {
68+
if !skip_implementation && def.stub_or_impl == FunctionStubOrImpl::Impl {
69+
self.error(
70+
errors,
71+
def.id_range,
72+
ErrorKind::InvalidOverload,
73+
None,
74+
"@overload decorator should not be used on function implementations.".to_owned(),
75+
);
76+
}
77+
6878
// This function is decorated with @overload. We should warn if this function is actually called anywhere.
6979
let successor = self.bindings().get(idx).successor;
7080
let ty = def.ty.clone();
7181
if successor.is_none() {
7282
// This is the last definition in the chain. We should produce an overload type.
7383
let mut acc = Vec1::new((def.id_range, ty));
74-
let mut first = def;
75-
while let Some(def) = self.step_overload_pred(predecessor) {
76-
acc.push((def.id_range, def.ty.clone()));
77-
first = def;
84+
let mut has_overload_after = false;
85+
let mut has_implementation_before_overload = false;
86+
let mut has_any_implementation = false;
87+
let mut temp_pred = *predecessor;
88+
while let Some(current_pred_idx) = temp_pred {
89+
let mut current_binding = self.bindings().get(current_pred_idx);
90+
while let Binding::Forward(forward_key) = current_binding {
91+
current_binding = self.bindings().get(*forward_key);
92+
}
93+
if let Binding::Function(func_idx, next_predecessor, _) = current_binding {
94+
let func_def = self.get_idx(*func_idx);
95+
96+
if func_def.metadata.flags.is_overload {
97+
has_overload_after = true;
98+
}
99+
if func_def.stub_or_impl == FunctionStubOrImpl::Impl {
100+
has_any_implementation = true;
101+
if !func_def.metadata.flags.is_overload {
102+
has_implementation_before_overload = true;
103+
}
104+
}
105+
if has_overload_after && has_any_implementation && has_implementation_before_overload {
106+
break;
107+
}
108+
temp_pred = *next_predecessor;
109+
} else {
110+
break;
111+
}
112+
}
113+
114+
let mut first = def.clone();
115+
while let Some(predecessor_def) = self.step_overload_pred(predecessor) {
116+
acc.push((predecessor_def.id_range, predecessor_def.ty.clone()));
117+
first = predecessor_def;
78118
}
79119
if !skip_implementation {
80-
self.error(
81-
errors,
82-
first.id_range,
83-
ErrorKind::InvalidOverload,
84-
None,
85-
"Overloaded function must have an implementation".to_owned(),
86-
);
120+
if !has_implementation_before_overload && def.stub_or_impl == FunctionStubOrImpl::Impl {
121+
self.error(
122+
errors,
123+
def.id_range,
124+
ErrorKind::InvalidOverload,
125+
None,
126+
"@overload decorator should not be used on function implementations.".to_owned(),
127+
);
128+
} else if has_any_implementation {
129+
self.error(
130+
errors,
131+
def.id_range,
132+
ErrorKind::InvalidOverload,
133+
None,
134+
"@overload declarations must come before function implementation. ".to_owned(),
135+
);
136+
}
137+
else {
138+
self.error(
139+
errors,
140+
first.id_range,
141+
ErrorKind::InvalidOverload,
142+
None,
143+
"Overloaded function must have an implementation".to_owned(),
144+
);
145+
}
87146
}
88-
if acc.len() == 1 {
147+
if acc.len() == 1 && !has_overload_after && !has_any_implementation {
89148
self.error(
90149
errors,
91150
first.id_range,
@@ -110,29 +169,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
110169
}
111170
} else {
112171
let mut acc = Vec::new();
113-
let implementation_range = def.id_range;
114172
let mut first = def;
115-
let mut has_overload_successor = false;
116-
117-
let binding = self.bindings().get(idx);
118-
let mut current_successor = binding.successor;
119-
while let Some(succ_key_idx) = current_successor {
120-
let succ_def = self.get_idx(succ_key_idx);
121-
if succ_def.metadata.flags.is_overload {
122-
has_overload_successor = true;
123-
break;
124-
}
125-
current_successor = self.bindings().get(succ_key_idx).successor;
126-
}
127-
if has_overload_successor && !skip_implementation {
128-
self.error(
129-
errors,
130-
implementation_range,
131-
ErrorKind::InvalidOverload,
132-
None,
133-
"Function implementation must come after all @overload declarations. ".to_owned(),
134-
);
135-
}
136173
while let Some(def) = self.step_overload_pred(predecessor) {
137174
acc.push((def.id_range, def.ty.clone()));
138175
first = def;
@@ -247,16 +284,6 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
247284
self_type = self_type.map(Type::type_form);
248285
}
249286

250-
if is_overload && stub_or_impl == FunctionStubOrImpl::Impl {
251-
self.error(
252-
errors,
253-
def.name.range,
254-
ErrorKind::InvalidOverload,
255-
None,
256-
"@overload decorator should not be used on function implementations.".to_owned(),
257-
);
258-
}
259-
260287
// Determine the type of the parameter based on its binding. Left is annotated parameter, right is unannotated
261288
let mut get_param_ty = |name: &Identifier, default: Option<&Expr>| {
262289
let ty = match self.bindings().get_function_param(name) {
@@ -495,6 +522,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
495522
id_range: def.name.range,
496523
ty,
497524
metadata,
525+
stub_or_impl,
498526
})
499527
}
500528

pyrefly/lib/alt/types/decorated_function.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use crate::types::callable::FuncId;
2121
use crate::types::callable::FuncMetadata;
2222
use crate::types::callable::FunctionKind;
2323
use crate::types::types::Type;
24+
use crate::binding::binding::FunctionStubOrImpl;
2425

2526
/// The type of a function definition after decorators are applied. Metadata arising from the
2627
/// decorators can be stored here. Note that the type might not be a function at all, since
@@ -30,6 +31,7 @@ pub struct DecoratedFunction {
3031
pub id_range: TextRange,
3132
pub ty: Type,
3233
pub metadata: FuncMetadata,
34+
pub stub_or_impl: FunctionStubOrImpl,
3335
}
3436

3537
impl Display for DecoratedFunction {
@@ -51,6 +53,7 @@ impl DecoratedFunction {
5153
})),
5254
flags: FuncFlags::default(),
5355
},
56+
stub_or_impl: FunctionStubOrImpl::Stub,
5457
}
5558
}
5659
}

pyrefly/lib/binding/binding.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -754,7 +754,7 @@ impl IsAsync {
754754
}
755755

756756
/// Is the body of this function stubbed out (contains nothing but `...`)?
757-
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
757+
#[derive(Clone, Copy, Debug, PartialEq, Eq, TypeEq, VisitMut)]
758758
pub enum FunctionStubOrImpl {
759759
/// The function body is `...`.
760760
Stub,

pyrefly/lib/test/overload.rs

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -412,13 +412,30 @@ testcase!(
412412
from typing import overload
413413
414414
@overload
415-
def f(x: int) -> int: ... # E: Overloaded function must have an implementation
415+
def f(x: int) -> int: ...
416416
417417
@overload
418-
def f(x: str) -> str: ...
418+
def f(x: int | str) -> int | str: # E: @overload decorator should not be used on function implementations.
419+
return x
419420
420421
@overload
421-
def f(x: int | str) -> int | str: # E: @overload decorator should not be used on function implementations.
422+
def f(x: str) -> str: ... # E: @overload declarations must come before function implementation.
423+
"#,
424+
);
425+
426+
427+
testcase!(
428+
test_implementation_before_overload,
429+
r#"
430+
from typing import overload
431+
432+
def f(x: int | str) -> int | str:
422433
return x
434+
435+
@overload
436+
def f(x: int) -> int: ...
437+
438+
@overload
439+
def f(x: str) -> str: ... # E: @overload declarations must come before function implementation.
423440
"#,
424441
);

0 commit comments

Comments
 (0)